Emergent temporal abstractions in autoregressive models enable hierarchical reinforcement learning | Summary

Contents

Summary
Key results
Results
Discussion
Appendix
Brief Thoughts

This article explains the key points of Emergent temporal abstractions in autoregressive models enable hierarchical reinforcement learning.

2025-12-23 (arXiv)
Kobayashi, Seijin, Schimpf, Yanick, Schlegel, Maximilian, Steger, Angelika, Wolczyk, Maciej, von Oswald, Johannes, Scherrer, Nino, Maile, Kaitlin, Lajoie, Guillaume, Richards, Blake A., et al.
Google, Paradigms of Intelligence Team, Google
Paper

Read this article in Korean

Summary

The paper investigates whether autoregressive next-action predictors develop reusable behaviors that extend across multiple timesteps and support learning from sparse rewards. Experiments use discrete gridworld navigation and continuous MuJoCo ant control, with post-training tasks requiring longer, previously unseen combinations of subgoals.
The method pretrains an autoregressive action model, freezes it, and trains a future-conditioned metacontroller to generate linear interventions in its residual stream. A learned switching gate maintains controller codes across multiple timesteps, while a variational objective encourages meaningful sampling from a standard normal prior without subgoal labels.
Internal RL replaces the future-conditioned code encoder with a causal policy and reinforces controller choices at learned switching events. It learns the tested sparse-reward tasks while raw-action RL, adapted CompILE, cotraining, and forced-switching variants fail to achieve comparable success; the experiments use small models and deliberately hierarchical environments.

Key results

The paper reports five results: next-action prediction develops representations of hidden subgoals; linear internal controllers steer behavior toward those subgoals; sequencing controllers enables length and compositional generalization; a metacontroller discovers controllers and termination conditions without subgoal labels; and internal RL composes these abstractions using terminal rewards.

Linear probing measures whether subgoal information is accessible, while interventions test whether controllers can causally steer behavior.
Unsupervised discovery reads from and writes to a frozen pretrained residual stream and uses future information during training.
Internal RL uses residual activations as observations and controller codes as actions, reducing action dimensionality and the number of policy decisions.

Residual-stream controller discovery and reinforcement learning over temporally extended internal actions.

Results

The experiments first examine subgoal representations and their controllability, then train a metacontroller without subgoal labels, and finally evaluate sparse-reward RL on unseen subgoal combinations.

Linearly controllable abstract action representations emerge in autoregressive models

Base models are trained from scratch on expert observation-action trajectories without rewards, task descriptors, or subgoal labels. A causal transformer models gridworld behavior, and a Hawk state-space model models ant control; training combines next-action likelihood with an auxiliary next-observation objective weighted by λ ≥ 0.

Discrete actions use a categorical likelihood; continuous actions use a Gaussian likelihood with learned diagonal covariance.
Tasks require visiting colored locations in order while avoiding walls and incorrect colored locations.
An abstract subgoal comprises two consecutive colored locations in gridworld and one colored location in ant control.

Hierarchical navigation environments and the composition of unseen post-training subgoal sequences.

Linear classifiers decode latent subgoals from instantaneous residual-stream activations in the gridworld model. In the illustrated six-layer transformer, predictions become more concentrated on the current subgoal as behavioral evidence accumulates, and decoding is stronger in middle and deeper layers than at the input.

The probes use ground-truth subgoal labels for analysis; those labels are absent from base-model pretraining.
Decodability establishes accessible subgoal information but does not alone establish behavioral control.

Linear-probe subgoal distributions across time and transformer depth.

For causal intervention, a linear controller applies \(e_{t,l}\leftarrow e_{t,l}+U_t e_{t,l}\) between model blocks. Controllers trained and selected using ground-truth subgoal labels steer the frozen base model through longer, unseen subgoal combinations; Figure 4 shows better control near the middle of the model across the tested backbone depths.

Results are averaged over 5 seeds, with no retraining of the autoregressive model.
Strong decoding in later layers and stronger control near mid-depth motivate mid-depth interventions in subsequent experiments.
This experiment supplies the subgoal sequence externally; autonomous controller selection is evaluated during RL.

Post-training compositional success by backbone depth and controller insertion layer, averaged over 5 seeds.

Unsupervised metacontroller discovers temporally-abstract actions within autoregressive models

The unsupervised metacontroller is a stochastic recurrent hypernetwork: its encoder proposes Gaussian controller codes, and its decoder maps each code to a linear controller matrix. During training, the encoder receives an embedding of the entire activation trajectory to infer goals that may be ambiguous from early actions alone; the switching unit remains causal.

The base model stays frozen, and metacontroller training uses cleaner demonstrations than base-model pretraining.
Future conditioning supports offline abstraction discovery and is removed from the action policy used during RL.
The decoder produces intervention parameters rather than motor actions.

Metacontroller modules for future-conditioned code inference, temporal integration, and linear residual-stream intervention.

Temporal integration follows z_t = β_t ⊙ z̃_t + (1 − β_t) ⊙ z_t−1, with \(\beta_t\in[0,1]\). Small gate values preserve the previous controller code, and values near one adopt a new proposal; next-action prediction and a KL penalty toward \(\mathcal N(0,I)\) train the metacontroller without explicit subgoal or boundary labels.

The KL weight α controls the information bottleneck and encourages meaningful controller sampling from the prior.
Learned gates become sparse and approximately binary, with switches aligned to subgoal changes in the illustrated ant trajectories.
The unconditional proposal prior permits sampling new controller codes independently of earlier proposals.

Learned switching gates aligned with colored subgoal segments in ant-control trajectories.

Rate-distortion experiments sweep α and compare action-prediction loss with KL divergence to the prior. Subgoal-aligned switching appears over a range of settings with a frozen pretrained base model; the tested cotraining variants lose this structure and often collapse to a single initial switch.

Figure 7 places aligned switching near a gap in the rate-distortion frontier, where a small increase in latent information substantially reduces prediction error.
These comparisons support freezing the pretrained model under the tested training procedures.
Appendix C.6 uses one randomly selected backbone and 10 metacontroller seeds per KL setting.

Rate-distortion comparisons between frozen-backbone and jointly trained metacontrollers in gridworld and ant control.

Internal reinforcement learning

Post-training provides reward 1 only after completion of the entire required subgoal sequence, and reward 0 otherwise. Raw-action RL fine-tunes the pretrained predictor with a critic-free, PPO-style algorithm related to GRPO, but output-level sampling yields too few successful novel combinations for substantial learning.

The paper describes successful raw-action sampling as occurring on the order of one in a million trajectories in the challenging setting.
The algorithm normalizes trajectory rewards across the batch to estimate advantages and omits GRPO’s grouping mechanism.
The raw-action baseline begins from behavioral pretraining.

Internal RL trains a causal recurrent policy over controller codes while fixing all other modules, including the base model, controller decoder, and switching mechanism. The switching gate is hardened with threshold 0.5, so a selected controller persists until a learned switching event; the fixed modules and physical environment together form the policy’s effective environment.

Controller codes have dimension 8, compared with residual-stream dimension 256.
A one-layer SSM policy observes residual activations and maintains recurrent state because a single activation vector may omit information from the input history.
Decisions at switching events shorten the effective policy horizon.

Figure 8 shows increasing success for full internal RL in both environments, while raw-action RL, adapted CompILE, cotraining, and forced switching remain near zero on the linear scale. Curves report medians and the 25th–75th quantile range over 30 runs, with 3 metacontroller seeds for each of 10 pretrained models.

The forced-switching variant uses \(\beta_t=1\) at every timestep during abstraction discovery as well as RL, testing the removal of temporal persistence throughout training.
Log-scale results show that some baselines obtain occasional successes but fail to reinforce them into successful policies.
Ant success in Figure 8 is normalized by the expert’s post-training success rate of 0.7; plotted ant scores are therefore not raw completion probabilities.

Sparse-reward learning curves for internal RL, raw-action RL, adapted CompILE, cotraining, and forced switching.

Discussion

The Discussion interprets the experiments as evidence that next-action prediction can develop internal representations of temporally extended, goal-directed behavior. The method extracts controllers and termination conditions from these representations and learns to compose them through RL.

The authors relate the approach to alternating self-supervised learning and RL and to architectures that modulate a predictive model for a selected goal.
The metacontroller maintains recurrent state and learns interventions that improve behavioral prediction, whereas the sparse autoencoders discussed here reconstruct instantaneous activations.

The authors caution that their experimental environments are controlled and that applications to larger models and real-world tasks require further investigation. LLM reasoning and large-scale model control are proposed directions, without language-model reasoning benchmarks or real-world robot evaluations in this paper.

Appendix

A. Environment details describes the two navigation environments. A.1. Gridworld-pinpad and A.1.1. Markov decision process specification define cardinal movement on a randomized grid, with one-hot observations of objects, walls, and agent position; moves into walls or outside the grid have no effect. Incorrect colored-cell visits and timeouts terminate episodes, and only full task completion earns reward 1. A.1.2. Task specification and hyperparameters sets a 7-by-7 grid, 8 colored cells, 4 walls, and maximum episode length 100. Abstract subgoals are 0–1, 2–3, 4–5, and 6–7; the post-training sequence is 0–1–2–3–4–5–6–7–0–1–2–3.
A.2. Ant-pinpad and A.2.1. Markov decision process specification define continuous navigation using 8-dimensional joint-torque actions. Layouts, initial position, yaw, joint angles, and velocities are randomized; observations combine symlog-transformed proprioception, normalized global coordinates, relative object and wall locations, and local coordinates. Wall entry, an incorrect colored-cell visit, timeout, or torso height outside [0.2, 1.0] terminates an episode, while full completion earns reward 1. A.2.2. Task specification and hyperparameters sets a 4-by-4 grid, 4 colored cells, 1 wall, maximum episode length 500, and post-training sequence 0–1–2–3.
B. Additional experimental results includes further probing, abstraction, and RL analyses. B.1. Belief state probing reports that gridworld subgoals become more linearly decodable around the middle of the six-layer transformer, with accuracy rising from 30% to about 50%; output-layer probing can deteriorate after 100K backbone steps. B.2. Effect of sequence model training hyperparameters on the abstract action representations finds benefits from longer pretraining, moderate weight decay, and an auxiliary observation objective, while excessive weight decay or observation-loss weighting reduces controller generalization. Base representations remain robust to the tested expert action noise; the authors suggest that failures at very low observation-loss weights may be environment-specific.
B.3. Unsupervised abstract action discovery and B.3.1. Temporal abstraction in the gridworld show that learned switching also aligns with gridworld subgoals. B.3.2. Quality of abstract actions tests codes collected from prior-sampled ant rollouts: injecting a “go to blue” code at timestep 30 increases goal-reaching success from 23% to 36%, and activating an object-1 code after object 0 increases success on that unseen transition from 10% to 24%. B.4. Internal reinforcement learning uses log-scale curves to distinguish occasional baseline successes from the sustained improvement achieved by full internal RL.
C. Experimental details specifies the training and evaluation procedures. C.1. Pretraining of sequence models and C.1.1. Expert trajectory generation use dynamic-programming shortest-path experts for gridworld and a PPO-trained expert for ant control. The ant expert receives additional directional observations and an intrinsic velocity-direction reward; its success is 0.8 on pretraining tasks and 0.7 on post-training tasks. C.1.2. Sequence model training combines next-action and next-observation likelihoods, and C.1.3. Seed specifies 10 independently pretrained models per environment. The listed backbones have embedding dimension 256: a six-layer transformer trained for 256000 steps in gridworld and an eight-layer SSM trained for 204800 steps in ant-pinpad.
C.2. Belief state probing trains supervised linear classifiers on frozen activations. C.3. Controller compositional generalization supplies one-hot subgoal identifiers at each timestep to isolate controllability, using rank-16 linear controllers in gridworld and linear controllers in ant-pinpad. C.4. Unsupervised abstract action discovery learns Gaussian controller codes under KL regularization without subgoal labels, with the listed configurations intervening at half the backbone depth and using code dimension 8.
C.4.1. Baseline – forced resets sets the switching gate to one at every timestep. C.4.2. Baseline – metacontroller cotraining jointly trains backbone and metacontroller, starting from random backbone parameters in gridworld and pretrained parameters in ant-pinpad. C.4.3. Baseline – CompILE adapts the latent-code module to consume raw observations and infer segment boundaries, with a maximum of 4 segments; the sequence model is initialized from pretraining. C.4.4. Metacontroller training dataset uses cleaner demonstrations, including mean ant-expert actions and zero gridworld action noise. C.4.5. Seed repeats each abstraction-discovery configuration over 3 seeds for each of the 10 sequence models.
C.5. RL experiments and C.5.1. Internal RL provide algorithms for effective-environment initialization, abstract-action execution, and policy training, using a one-layer SSM policy with embedding dimension 256. C.5.2. RL algorithm details uses a clipped PPO-style surrogate and critic-free advantages computed by normalizing trajectory rewards across the batch. C.5.3. Baseline – raw action RL fine-tunes the pretrained backbone, while C.5.4. Baseline – others applies internal RL to alternative abstraction models. C.5.5. Seed and Hyperparameter selection searches learning rates (0.000003, 0.00001, 0.00003, 0.0001, 0.0003) and selects configurations by median RL performance across 10 pretrained models and 3 seeds.
C.6. Rate-distortion curve uses 10 abstraction-discovery seeds per KL setting with one randomly selected backbone and a Gumbel-Sigmoid switching variant for gridworld. Switches are counted after thresholding at 0.5, and patterns are manually checked for alignment with subgoal changes. The analysis uses fewer backbone seeds than the main RL comparison.
D. Architecture details and D.1. Sequence model specify pre-normalized residual architectures. D.1.1. SSM uses Hawk recurrence, and D.1.2. Transformer uses attention with relative-position biases; both include MLP channel mixing. D.2. Metacontroller architecture combines a GRU history state, a full-trajectory sequence embedding, a Gaussian code encoder, a causal switching unit, temporal integration, and a hypernetwork decoder for linear interventions. D.3. Internal RL policy architecture uses a recurrent one-layer SSM because a single residual activation may omit information from the observation history.
E. Additional discussions includes the latent-variable formulation and a comparison of RL gradient estimators. E.1. Graphical model and ELBO derivation describes a process that retains its controller code between switches and draws a new code from a standard normal prior at a switch. Training relaxes switching continuously and assumes causal residual activations contain enough information to emit termination signals without future conditioning. E.2. Internal RL vs RL with reparametrization trick compares estimators under terminal rewards and a fixed-variance abstract policy, further assuming one abstract decision for the variance comparison. The analysis describes a potential reduction in horizon-dependent noise alongside a tradeoff in variance over latent-code sampling.

Brief Thoughts

Causal controller interventions, subgoal-aligned switching, and RL ablations together connect hidden subgoal information to reusable behavior. The forced-switching comparison shows that, in these tasks, temporal persistence contributes beyond exploration through latent-code noise alone.

The method depends on behavioral pretraining that supports useful internal controllers and on clean demonstrations for abstraction discovery. Its effectiveness is established for explicitly compositional navigation tasks; transfer to less structured behavior, unreliable termination signals, and large language models remains an open empirical question.

Gorio Tech Blog