Pretraining and meta-training: giving a model a fish versus teaching it how to fish
Training originally began with the goal of teaching models to complete a wide variety of tasks. In large-scale pretraining, beyond memorizing the training tasks themselves, the most important emergent capability is in-context learning. This capability enables large models to construct solutions to new tasks from information provided in a prompt.
OmniRL proposes a large-scale meta-training method. Unlike pretraining, OmniRL does not aim to memorize skills for specific tasks. Its goal is to learn the process of “how to perform reinforcement learning” itself. Meta-learning, or “learning to learn,” was proposed as early as the 1980s. The OmniRL paper argues, however, that meta-learning without large-scale task diversity and long sequences easily collapses into a “task identification” pattern: the model merely memorizes its training environments and, during inference, identifies the current environment to activate the corresponding skill. This pattern cannot generalize to unseen or out-of-distribution tasks.
Random worlds: AnyMDP
An example random world generated by AnyMDP. Node color represents the average reward of a state, and edge intensity represents the average transition probability between states.
AnyMDP constructs random transition probabilities and reward functions from Markov decision processes (MDPs). It can quickly and inexpensively generate a massive, scalable collection of environments for meta-reinforcement learning. We generated more than 500,000 distinct tasks and synthesized over 10 billion time steps of meta-training data from them. Individual sequences reached lengths of more than one million time steps.
Unifying multiple forms of reinforcement learning and imitation learning through in-context learning for the first time
OmniRL uses both prior information and posterior reward feedback for in-context learning, enabling a model to switch autonomously among learning modes as needed. The examples below show an OmniRL model trained in random worlds. Using only in-context learning, without any gradient-based optimization, the model performs well from a cold start or from a supplied demonstration trajectory, whether expert or suboptimal. It autonomously switches among online reinforcement learning, offline reinforcement learning, and imitation learning, demonstrating the remarkable flexibility of in-context learning. The model can also continue to improve through independent exploration after receiving a demonstration.
OmniRL’s performance in completely unseen Gymnasium environments.
Agents trained by OmniRL can even complete cooperative multi-agent tasks. By including the other agent’s state in the observation space, the model can solve simple tasks such as Switch, which require agents to exhibit different behavior patterns in order to cooperate. Through in-context learning and adaptation, two OmniRL-controlled agents can collaborate effectively to complete these tasks.
Revealing the underlying importance of data diversity and sequence length for the first time
The relationship between positional loss, meta-training steps, and context length.
OmniRL uses Transformers with tens of millions of parameters and efficient linear-attention architectures. It trains on more than 500,000 tasks with sequence lengths exceeding one million time steps. Experiments compare equal amounts of data drawn from different numbers of tasks. When the number of tasks is insufficient, the model shifts toward a “memorization plus environment identification” strategy: it stores every training environment in parametric memory and identifies the current one quickly from context. An agent in this regime can adapt to a previously seen environment with fewer examples, but it cannot generalize to an unseen environment. Only a sufficiently large task set elicits genuine, general-purpose in-context learning. This capability generalizes effectively to unseen tasks, but it requires a longer in-context learning period for every task.
These findings suggest that:
- Completeness and diversity of data matter more than absolute fidelity. Even distorted data may generalize better to real tasks if it improves the generality of in-context learning.
- Long-sequence modeling and long-term memory are a natural choice for general learning capabilities. As the number of training tasks increases, the model naturally stops memorizing task-specific knowledge and retains only the learning method. Consequently, even tasks present in the training set require longer adaptation. This is a defining property of large-scale meta-learning.
Clear advantages of linear self-attention in efficiency and long-sequence performance
The relationship between positional loss, meta-training steps, and context length.
OmniRL demonstrates the advantages of linear attention for the first time. As problem scale grows, the required context length grows with it, making the Transformer’s efficiency bottleneck increasingly severe. Linear attention, by contrast, offers clear advantages in both efficiency and long-sequence modeling. It also substantially outperforms sliding-window attention over long temporal spans, showing that AnyMDP provides an excellent evaluation environment for long sequences.
Technical exploration toward the next generation of general embodied agents
Our ultimate goal is an agent that can explore and learn fully autonomously in any environment, a capability with particular importance for embodied intelligence. Large language models capture extensive common sense, encyclopedic knowledge, mathematics, and logic in parametric memory, creating the foundation for their zero-shot capabilities. Embodied intelligence, however, faces diverse environments and tasks as well as complex variation among physical embodiments. Common sense alone is unlikely to provide the core solution. We believe autonomous learning and long-term memory will be essential to general embodied agents.
Similarities and differences from long-horizon reasoning and chain-of-thought in current large language models
OmniRL currently focuses more on the learning capabilities of System 1, or intuitive thought, while long-horizon reasoning and chain-of-thought in large language models focus more on System 2 itself, including logical reasoning and planning. Mainstream large models have not yet sufficiently explored how either System 1 or System 2 capabilities are learned and improved. OmniRL fills many of these gaps.