Emergent temporal abstractions in autoregressive models enable hierarchical reinforcement learning

TL;DR

This paper introduces a hierarchical reinforcement learning approach using autoregressive models' internal representations, enabling efficient exploration and learning from sparse rewards by controlling internal activations with higher-order sequence models.

Contribution

It presents a novel higher-order, non-causal sequence model that controls autoregressive model activations, facilitating hierarchical RL and internal reinforcement learning within foundation models.

Findings

Higher-order models learn to compress long activation sequences.

Controllers execute long-timescale, meaningful actions with learned termination.

Internal RL improves learning from sparse rewards where standard RL fails.

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.