Emergent temporal abstractions in autoregressive models enable hierarchical reinforcement learning

Results

• The study pretrains action models on unlabeled behavior from hierarchical grid-world and MuJoCo ant tasks, then evaluates them on longer tasks with new subgoal orders.The behavioral sequences contain no reward or subgoal labels.Source: Linearly controllable abstract action… · p. 3, Figure 2 · p. 3
• Linear probes decode subgoals from residual activations more accurately with accumulated evidence and greater depth, peaking near the final transformer embedding.The figure reports increasing decoding performance through layer 4, remaining strong through the final embedding at layer 6.Source: Linearly controllable abstract action… · p. 3, Figure 3 · p. 4
• Mid-depth controllers solve longer tasks with unseen subgoal orders by activating ground-truth-selected subgoal controllers in sequence, without retraining the autoregressive model.The post-training tasks test both length and compositional generalization.Source: Linearly controllable abstract action… · p. 4, Figure 4 · p. 4

Unsupervised metacontroller discovers temporally-abstract actions within autoregressive models

• A non-causal encoder proposes controller codes from the full activation sequence, while a switching unit integrates new codes at a learned, time-varying rate.Self-supervised learning combines next-action prediction with prior matching to encourage meaningful code sequences and sparse, subgoal-aligned switching.Source: Unsupervised metacontroller discovers… · p. 5, Figure 5 · p. 5
• The metacontroller recovers ground-truth subgoal switching times, with quasi-binary sparse switching and essentially perfect temporal segmentation in the ant environment.Its controller codes correspond to meaningful abstract actions and generalize to new task configurations and switching times.Source: Unsupervised metacontroller discovers… · pp. 5–6, Figure 6 · p. 5
• With a frozen base model, subgoal-aligned switching appears across a range of rate-distortion trade-offs; co-training instead loses this structure and usually converges to one initial switch.The authors conclude that controlling a frozen pretrained action predictor is essential for discovering temporal abstractions.Source: Unsupervised metacontroller discovers… · p. 6, Figure 7 · p. 6
• Internal RL achieves high success on post-training tasks, while disabling temporal integration, co-training the base model, and CompILE yield near-zero success.The tasks require sparse-reward hierarchical behavior; the full method combines high initial success with efficient credit assignment.Source: Figure 8 · p. 7, Unsupervised metacontroller discovers… · p. 7

Discussion

• Internal controllers first use ground-truth subgoals to solve complex tasks compositionally; a metacontroller later selects temporally abstract actions without those labels.Direct reinforcement of the internal activation controllers enables learning in hierarchical sparse-reward tasks where other RL techniques fail.Source: Discussion · pp. 7–8
• The work adds evidence to debate over whether autoregressive predictors form temporal abstractions and plans, using tasks requiring behavior recombination, sparse-reward learning, and imitation of goal-directed agents.The authors describe solving sparse-reward tasks without manual reward shaping as a critical step toward autonomous agents in open-ended search spaces.Source: Discussion · p. 8
• The reported advantage over standard RL methods warrants care because the experiments use controlled environments, and adaptation to larger models and tasks remains future work.The environments are challenging enough for GRPO and CompILE to fail, despite their simplicity.Source: Discussion · p. 8
• In continuous environments with hidden discrete hierarchical structure, the models discover that structure and the metacontroller develops sparse, quasi-binary switching units.The authors note that the models use continuous units and are trained by gradient descent.Source: Discussion · p. 8
• Unlike JEPA’s proposal to learn abstract representations without autoregressive prediction, this approach uses raw action prediction, which partly enables discovering task decompositions into subgoals.Task decomposition is identified as an open problem in the JEPA proposal.Source: Discussion · p. 8
• Metacontrollers are predictive, intervene on the residual stream, and maintain internal state, unlike instantaneous sparse autoencoders; the authors see potential for interpretable interventions over extended periods.They are trained to lower next-token prediction error, and the authors propose investigating their capabilities in larger-scale models such as LLMs.Source: Discussion · pp. 8–9

A.1. Gridworld-pinpad

• The gridworld-pinpad environment places an agent among uniquely colored cells and requires visiting a task-specific sequence of them.Source: §A.1 · p. 12

• Tasks specify an ordered sequence of colored cells to visit, and actions correspond to the four cardinal directions.Source: §A.1.1 · p. 12
• The state is a G-by-G grid with O colored cells and W walls, the agent occupying a non-wall cell, plus a record of cells visited during the episode.Source: §A.1.1 · p. 12
• Moves toward a wall or outside the grid are no-ops; an episode ends on successful completion, an out-of-order colored-cell visit, or exceeding T steps.A cell counts as visited when the agent moves onto it from a different cell.Source: §A.1.1 · p. 12
• At episode start, the colored cells, walls, and agent position are sampled randomly without overlap.Source: §A.1.1 · p. 12
• Observations one-hot encode the contents of each cell and the agent’s position, and successful task completion yields reward 1 while other outcomes yield 0.Source: §A.1.1 · p. 12

A.1.2. Task specification and hyperparameters

• Gridworld-pinpad uses G=7, O=8, W=4, and T=100 for both pretraining and post-training tasks.Source: §A.1.2 · p. 12
• The gridworld pretraining tasks combine abstract subgoals 0–1, 2–3, 4–5, and 6–7, while the post-training task is 0–1–2–3–4–5–6–7–0–1–2–3.Source: §A.1.2 · p. 12

A.2. Ant-pinpad

• Ant-pinpad is the continuous-control counterpart of gridworld-pinpad, with a MuJoCo ant visiting a task-specific sequence of colored cells.Source: §A.2 · p. 12
• Table A2 lists the ant-pinpad pretraining tasks, with each entry denoting a sequence of colors to visit in order.Source: Table A1 · p. 12

• Tasks specify a sequence of colored cells to visit, and the state includes the grid, colored cells, walls, ant proprioception and position, and cells already visited.Source: §A.2.1 · p. 12
• The action is an eight-dimensional continuous vector of torques applied to the ant’s eight joints.Source: §A.2.1 · p. 12
• Episodes end if the ant enters a wall, its torso leaves the 0.2–1.0 vertical range, the task is completed, a wrong colored cell is visited, or T timesteps elapse.A colored cell counts as visited when the ant enters it from a different cell.Source: §A.2.1 · p. 12
• Initial layouts and agent positions are randomized; the ant is placed within the sampled starting cell with random yaw and small random joint-angle and velocity offsets.Source: §A.2.1 · p. 12
• Observations combine symlog-transformed proprioception with normalized global coordinates, relative cell and wall positions, and the ant’s local cell coordinates; successful completion yields reward 1, otherwise 0.Source: §A.2.1 · p. 13

• Both pretraining and post-training use G=4, O=4, W=1, and T=500.Source: §A.2.2 · p. 13
• The post-training task is 0−1−2−3; the pretraining tasks are listed in Table A2.Source: §A.2.2 · p. 13

• Probe accuracy rises from 30% to about 50% around the middle of the six-layer model, although subgoals are not explicitly encoded in the input data.Source: §B.1 · p. 13, Figure A1 · p. 13
• Near the output layer, probe performance deteriorates for backbones trained beyond 100K steps.Source: §B.1 · p. 13, Figure A1 · p. 13

B.2. Effect of sequence model training hyperparame-

• Longer sequence-model training generally improves controller generalization across base-model depths 4, 6, and 8.Source: §B.2 · p. 13, Figure A2 · p. 14
• Weight decay benefits internal representations across model depths, but too much degrades them, indicating a regularization trade-off.Source: §B.2 · p. 13, Figure A2 · p. 14
• Some auxiliary next-observation loss improves abstract-action representations; very low coefficients sometimes prevent learning, a behavior authors suspect is environment-specific.Source: §B.2 · p. 13, Figure A2 · p. 14
• The abstract-action representation is robust to demonstrations with random non-terminating actions injected at probability 𝜖.Source: §B.2 · p. 14

B.3.1. Temporal abstraction in the gridworld

• In gridworld, the metacontroller’s quasi-binary switching gate aligns with the ground-truth abstract actions.Source: §B.3.1 · p. 14, Figure A3 · p. 14

B.3.2. Quality of abstract actions

• The learned quasi-discrete switches align with ground-truth subgoals, and controller codes form compact, steering-capable representations for internal RL.Source: §B.3.2 · pp. 14–15
• Injecting a “go to blue” code raises goal-reaching success from 23% to 36%, including when activated mid-rollout in new configurations and at unseen times.Source: §B.3.2 · p. 14, Figure A4 · p. 16
• For an untrained object-0-to-object-1 transition, activating the object-1 code raises success from 10% to 24%; codes were collected when the ant visited that object first.Source: §B.3.2 · pp. 14–15

B.4. Internal reinforcement learning

• Internal RL starts with the highest success rate and continues optimizing reward, while baselines fail to exploit their early successes.Source: §B.4 · p. 15, Figure A5 · p. 16

C.1. Pretraining of sequence models

• Base sequence models are pretrained to predict next actions and observations from observation histories, then held fixed for controller experiments.Source: §C.1 · p. 15, §C.1.2 · p. 15

C.1.1. Expert trajectory generation

• Gridworld demonstrations use a shortest-path policy with optional random actions; ant demonstrations come from PPO-trained agents, and expert success is imperfect.Source: §C.1.1 · p. 15

C.4. Unsupervised abstract action discovery

• Unsupervised discovery trains only the metacontroller on a frozen base model, using a regularized likelihood objective and a prior-compatible latent-code posterior.Source: §C.4 · p. 17

C.4.1. Baseline – forced resets

• The baselines force a switch at every step, jointly train the metacontroller and base model, or adapt CompILE to generate controller codes.Source: §C.4.1 · pp. 17–18, §C.4.2 · p. 18, §C.4.3 · p. 18

C.4.4. Metacontroller training dataset

• Cleaner expert demonstrations improve the metacontroller’s ability to compress abstract actions, so ant-pinpad uses expert mean actions and grid uses ε=0 demonstrations.The sequence model’s emergence of abstract actions is robust to suboptimality, but metacontroller compression improves as demonstrations more closely resemble the optimal stochastic policy.Source: §C.4.4 · pp. 18–19

C.4.5. Seed

• Unsupervised abstract-action discovery is run with three seeds for each of 10 pretrained sequence models and each hyperparameter configuration.Source: §C.4.5 · p. 19

C.5.1. Internal RL

• Internal RL learns over discovered abstract actions by treating the autoregressive model and part of the metacontroller as the environment.The RL algorithm is GRPO-related, modifies PPO, and replaces the learned critic with empirical advantage estimation.Source: §C.5.1 · p. 19, §C.5.2 · pp. 19–20

C.5.3. Baseline – raw action RL

• Raw-action RL finetunes the pretrained model without a metacontroller, while other baselines use internal RL with their discovered metacontrollers.Source: §C.5.3 · p. 20, §C.5.4 · p. 20

C.5.5. Seed and Hyperparameter selection

• Learning rates and hyperparameters are selected by the best median RL performance across 10 pretrained models and three seeds.Source: §C.5.5 · p. 20

C.6. Rate-distortion curve

• The rate-distortion comparison uses one randomly chosen sequence model, 10 seeds per KL strength, and compares the method with metacontroller cotraining.For gridworld-pinpad, it uses Gumbel-Sigmoid switching; switching sums are thresholded at 0.5 and checked for alignment with subgoal changes.Source: §C.6 · p. 20

D.1. Sequence model

• The base sequence model uses either transformer layers or Hawk state-space-model layers.Hawk is chosen for simplicity and computational efficiency; both architectures use pre-normalization, sequence mixing, and MLP channel mixing.Source: §D.1 · p. 20, §D.1.1 · p. 20, §D.1.2 · p. 20

• The metacontroller acts inside a frozen sequence-model backbone, emitting linear controllers that alter residual activations and can turn the model into a subgoal-optimizing policy.Source: §D.2 · pp. 20–21
• The latent subgoal at a timestep can require sequence-level information because its intent may become identifiable only as the trajectory unfolds.The metacontroller therefore uses an acausal sequence embedding alongside a recurrent history summary.Source: §D.2 · p. 21
• The metacontroller composes latent proposals sparsely in time, with β_t≈1 switching to a new subgoal and β_t≈0 retaining the previous one.The switching signal is causal, allowing switches to be emitted at test time without future information.Source: §D.2 · pp. 21–22
• Integrated latent codes are decoded into instantaneous controllers, and holding a code with β=0 lets the same controller act across multiple raw-action timesteps.Source: §D.2 · pp. 21–22

D.3. Internal RL policy architecture

• The internal RL policy uses a one-layer SSM because a single-layer residual activation may not contain the full raw-input history.Source: §D.3 · p. 22

E.1. Graphical model and ELBO derivation

• In the graphical model, β_t indicates a change in abstract action, and when β_t=0 the latent z_t remains equal to z_{t−1}.Source: §E.1 · p. 22
• The unsupervised objective maximizes an evidence lower bound on the likelihood of the observed action sequence.The variational distribution q is factorized according to the graphical model.Source: §E.1 · p. 22
• The experiments relax the sampled switching variable continuously, using its probability as the latent variable for training stability.Source: §E.1 · p. 22
• Experiments approximate the switching distribution using only past residual activations, assuming those activations are informative about when switches should occur.This permits causal switching at internal-RL test time; future conditioning could instead be retained and switches distilled to an unconditioned module.Source: §E.1 · p. 23

• The comparison assumes outcome supervision with one reward at the final timestep, and its variance analysis further considers a single abstract action drawn at sequence start.Source: §E.2 · p. 23
• PGraw has more structured expectation than PGz, and its variance with respect to ε can be zero while PGz variance scales with ε’s dimension.Source: §E.2 · p. 23
• PGraw’s variance can grow with timestep count and raw-action dimension, whereas PGz has O(1) variance, potentially improving credit assignment for compact, long-horizon abstract actions.Source: §E.2 · p. 23