Emergent temporal abstractions in autoregressive models enable hierarchical reinforcement learning

TL;DR

Addresses inefficiency of token-by-token RL exploration in pretrained autoregressive models under sparse rewards. Proposes internal RL by subsuming the base model into the environment and steering the residual stream with a metacontroller that discovers temporally-abstract actions as internal controllers. Introduces a future-conditioned, unsupervised metacontroller that generates controller codes and a switching mechanism to compose actions over long horizons. Validates on gridworld and MuJoCo ant tasks, showing latent actions enable compositional generalization, faster credit assignment, and successful learning where standard RL finetuning fails. Demonstrates that freezing the base model during metacontroller training yields favorable rate-distortion properties, supporting the discovery of robust internal abstractions and efficient hierarchical RL.

Abstract

Large-scale autoregressive models pretrained on next-token prediction and finetuned with reinforcement learning (RL) have achieved unprecedented success on many problem domains. During RL, these models explore by generating new outputs, one token at a time. However, sampling actions token-by-token can result in highly inefficient learning, particularly when rewards are sparse. Here, we show that it is possible to overcome this problem by acting and exploring within the internal representations of an autoregressive model. Specifically, to discover temporally-abstract actions, we introduce a higher-order, non-causal sequence model whose outputs control the residual stream activations of a base autoregressive model. On grid world and MuJoCo-based tasks with hierarchical structure, we find that the higher-order model learns to compress long activation sequence chunks onto internal controllers. Critically, each controller executes a sequence of behaviorally meaningful actions that unfold over long timescales and are accompanied with a learned termination condition, such that composing multiple controllers over time leads to efficient exploration on novel tasks. We show that direct internal controller reinforcement, a process we term "internal RL", enables learning from sparse rewards in cases where standard RL finetuning fails. Our results demonstrate the benefits of latent action generation and reinforcement in autoregressive models, suggesting internal RL as a promising avenue for realizing hierarchical RL within foundation models.

Figures (14)

• Figure 1: Research overview. (a) We let a metacontroller steer the residual stream activations of a pretrained autoregressive model. Through self-supervised next-action prediction, the metacontroller discovers how to generate sequences of simple (linear) internal controllers that change sparsely in time, following a dynamic switching unit $\beta_t \in [0,1]$. In hierarchically-structured tasks, each internal controller corresponds to a temporally-abstract action that leads the base autoregressive model to achieve a meaningful elementary goal. (b) We perform RL internally -- in the abstract space discovered by the metacontroller -- by subsuming the autoregressive model into the environment and acting in the residual stream on a contracted timescale.
• Figure 2: Environment and task design. (a) To complete a task, an agent must visit in sequence a number of subgoal locations, each marked with a specific color. The tasks are performed either in a discrete grid world or in a continuous motor control environment, illustrated above, where a quadrupedal robot (the 'ant') must be actuated at joint level. A task can be described as an abstract action sequence (the subgoal locations that must be visited), or as a sequence of low-level motor commands. (b) We pretrain autoregressive action models and metacontrollers on unlabeled behavioral datasets containing observation-action sequences of expert agents performing different tasks. These sequences do not contain rewards or subgoal labels. We then test the ability of the models to learn with RL tasks that comprise longer subgoal sequences, combined in new orders not seen during pretraining and metacontroller training.
• Figure 3: Internal belief distributions over abstract actions, according to a linear probe. Decoding performance of linear classifiers trained to predict groundtruth abstract actions from instantaneous residual stream activation vectors increases until mid-depth (layer 4) and remains strong up until the final embedding vector $e_{t,L}$ (here $L=6)$.
• Figure 4: Mid-depth linear internal controllers achieve length and compositional generalization. Both panels analyze success rate (the fraction of rewarded trials in which the full sequence of elementary goals defining a given task is completed) as a function of base model depth (the number of autoregressive model layers) and controlled layer (the layer at which the internal controller is inserted, with 0 corresponding to the middle of the base model). In both grid world (a) and ant (b) environments, inserting the controller near the middle layer results in better controllability, as measured by the success rate on the post-training tasks, which require both length and compositional generalization. To produce this analysis, we trained one controller per subgoal using groundtruth labels; to evaluate success rates we activated the controllers in correct order, again using groundtruth subgoal labels. Results averaged over 5 seeds.
• Figure 5: Details of the metacontroller architecture and the different modules at play. The metacontroller learns in a self-supervised way to generate sequences of internal controllers. Candidate controller codes $\tilde{z}_t$ are sampled from a Gaussian with context-dependent mean and covariance, and are integrated at a continuous, time-varying rate $\beta_t$, dynamically determined by a switching unit. Values of $\beta_t$ close to zero ignore new controller candidates; conversely, values close to unity lead to switching to a new controller. This mechanism is key for achieving temporal abstraction. The metacontroller features another key design element, a future-conditioned encoder: during self-supervised learning, the metacontroller is non-causal, and has access to the entire sequence of residual stream activations through a sequence embedding $s(e_{1:T})$.