Scaling of Search and Learning: A Roadmap to Reproduce o1 from Reinforcement Learning Perspective — part 1/8

Cleaned-text code points [0, 15697). Bibliography omitted here; source.md includes it.

# Scaling of Search and Learning: A Roadmap to Reproduce o1 from Reinforcement Learning Perspective

###### Abstract

OpenAI o1 represents a significant milestone in Artificial Inteiligence, which achieves expert-level performances on many challanging tasks that require strong reasoning ability. OpenAI has claimed that the main techinique behinds o1 is the reinforcement learining ( OpenAI, 2024a ; OpenAI, 2024b ) . Recent works use alternative approaches like knowledge distillation to imitate o1’s reasoning style, but their effectiveness is limited by the capability ceiling of the teacher model. Therefore, this paper analyzes the roadmap to achieving o1 from the perspective of reinforcement learning, focusing on four key components: policy initialization, reward design, search, and learning. Policy initialization enables models to develop human-like reasoning behaviors, equipping them with the ability to effectively explore solution spaces for complex problems. Reward design provides dense and effective signals via reward shaping or reward modeling, which is the guidance for both search and learning. Search plays a crucial role in generating high-quality solutions during both training and testing phases, which can produce better solutions with more computation. Learning utilizes the data generated by search for improving policy, which can achieve the better performance with more parameters and more searched data. Existing open-source projects that attempt to reproduce o1 can be seem as a part or a variant of our roadmap. Collectively, these components underscore how learning and search drive o1’s advancement, making meaningful contributions to the development of LLM.

One thing that should be learned from the bitter lesson is the great power of general purpose methods, of methods that continue to scale with increased computation even as the available computation becomes very great. The two methods that seem to scale arbitrarily in this way are search and learning . — Richard Sutton, 2019

## 1 Introduction

The field of Artificial Intelligence (AI) has witnessed unprecedented exploration and advancement of Large Language Models (LLMs) over the past two years. LLMs have progressively evolved to handle increasingly sophisticated tasks such as programming and solving advanced mathematical problems. OpenAI o1 represents a significant milestone in AI, which can generate very long reasoning process and conduct human-like reasoning actions like clarifying and decomposing questions, reflecting and correcting previous mistakes, exploring new solutions when encountering failure modes. The o1 model dramatically transcended the reasoning capabilities of preceding LLMs, achieving performance comparable to PhD-level proficiency. Its remarkable reasoning achievements signify OpenAI’s progression to the second stage (“Reasoner”) in their five-stage roadmap to Artificial General Intelligence (AGI).

The blog and system card of o1 demonstrate that the performance of o1 consistently improves with increasing the computation of reinforcement learning and inference ( OpenAI, 2024a ; OpenAI, 2024b ) . This suggests that o1 could drive two paradigm shifts in AI: from (self-)supervised learning toward reinforcement learning, and from scaling solely training computation to scaling both training and inference computation.

o1 scales up the train-time compute with reinforcement learning and the test-time compute with more thinking. We take search as a way to implement the thinking process of o1, since search is scalable ( Sutton, 2019 ) and there are many successful researches that use search for training and decision in reinforcement learning, like AlphaGo ( Silver et al., 2016 ) and AlphaGo Zero ( Silver et al., 2017 ) . In this paper, we take reinforcement learning as the core in the roadmap to o1. Our roadmap is illustrated in Figure 1 and consists of four components: policy initialization, reward design, search, and learning. We believe these four parts are the keys to constructing LLM with strong reasoning abilities like o1.

As depicted in Figure 2 , our roadmap starts with policy initialization. In the context of LLMs, the policy ( π ⁡ ( a | s ) \pi(a|s) ) typically refers to the probability distribution for generating the next token/step/response (action) based on a given context (state). Policy initialization brings human-like reasoning behaviors to LLMs, like task composition, self-evaluation and self-correction. Next, we enter the reward design, which aims to provide guiding signals for search and learning. The reward design can take or reshape the reward signal from the environment or learn a reward model from preference data. Both policy initialization and reward designs are preparations for search and learning. Search plays an important role in generating high-quality solutions at both training and testing phases, which can produce better solutions with more computation. Learning utilizes the data generated by search for improving the policy. The data used for learning is derived from the interaction of the LLM with the environment, rather than being manually curated by human experts, therefore eliminating the need for costly data annotation and enabling the potential for achieving superhuman performance.

##### Policy Initialization

Training an LLM from scratch using reinforcement learning is exceptionally challenging due to its vast action space. Fortunately, we can leverage extensive internet data to pre-train a language model, establishing a potent initial policy model capable of generating fluent language outputs. Moreover, prompt engineering and supervised fine-tuning help models acquire human-like reasoning behaviors, enabling them to think systematically and validate their own results. These approaches enable models to thoroughly explore their solution spaces, leading to more comprehensive problem-solving capabilities.

##### Reward Design

Both search and learning require guidance from reward signals to improve the policy. There are different levels of action granularity, each corresponding to varing levels of reward signals granularity, which can be explored further. Additionally, these signals are often sparse or even nonexistent in many environments. To transform sparse outcome reward to dense process reward, there are some reward shaping methods ( Ng et al., 1999 ) . For the environment where the reward signal is unavailable, like the task of story writing, we can learn a reward model from preference data ( Bai et al., 2022a ) or expert data ( Ng & Russell, 2000 ) .The construction of reward model can further evolve into building a world model ( Dawid & LeCun, 2023 ) .

##### Search

Search plays a crucial role during both the training and testing phases. The training time search refers to generating training data from search process. The advantage of using search to generate training data, as opposed to simple sampling, is that search yields better actions or solutions—i.e., higher-quality training data—thereby enhancing learning effectiveness. During inference, search continues to play a vital role in improving the model’s sub-optimal policies. For instance, AlphaGo ( Wan et al., 2024 ) employs Monte Carlo Tree Search (MCTS) during testing to enhance its performance. However, scaling test-time search may lead to inverse scaling due to distribution shift: the policy, reward, and value models are trained on one distribution but evaluated on a different one ( Gao et al., 2023 ) .

##### Learning

Learning from human-expert data requires costly data annotation. In contrast, reinforcement learning learns through interactions with the environment, eliminating the need for expensive data annotation and offering the potential for superhuman performance. In this roadmap, reinforcement learning utilizes data generated by search for learning via policy gradient or behavior cloning. Policy gradient methods have high data utilization, as they leverage both positive and negative solutions, whereas behavior cloning is advantageous in terms of simplicity and memory efficiency. A prominent example of the iterative interaction between search and learning is AlphaGo Zero ( Silver et al., 2017 ) , which combines Monte Carlo Tree Search (MCTS) ( Metropolis & Ulam, 1949 ) as the search algorithm with behavior cloning as the learning method, ultimately achieving superhuman performance in the game of Go.

We provide a detailed exploration of the potential implementations of Policy Initialization (Section 3 ), Reward Design (Section 4 ), Search (Section 5 ), and Learning (Section 6 ). Additionally, we review existing open-source o1 projects, illustrating how they may either serve as components of our framework or as specific instances within it (Section 7 ). Finally, we discuss the future development trends of o1 and the associated challenges (Section 8 ).

## 2 Background

Since this roadmap is designed from perspective of reinforcement learning, we introduce some background of reinforcement learning and its connection to LLM in this section. Unlike other learning paradigms, reinforcement learning learns through interaction with the environment, rather than learning from a static training dataset. In reinforcement learning, an agent learns by receiving rewards from the environment as it explores. Figure 3 illustrates the interaction between agent and environment in reinforcement learning for LLM.

##### Agent

Agent is the entity that interacts with environment, which makes decision according to its policy. Formally, a policy π \pi is a mapping from states to actions. It is often represented as a probability distribution ( π ⁡ ( a | s ) \pi(a|s) ) over actions given a state s s , where the agent selects actions based on these probabilities.

In the context of LLMs, an agent refers to LLM itself, its policy specify the probability distribution of either token-, step-, or solution-level actions based on the current state. The state s t s_{t} consists of the input provided to the model at time t t , including both user inputs and the model’s earlier outputs. The action taken by the model can vary depending on the problem setting; it involves generating a single token, completing a step, or providing a solution.

##### Environment

Environment refers to the system or world outside the agent. It responds to agent’s actions and provide feedback in terms of next state s t + 1 s_{t+1} and rewards r ⁡ ( s t , a t ) r(s_{t},a_{t}) .

Environmental feedback can be categorized as either deterministic or stochastic. Stochastic feedback is characterized by a transition distribution p ( s t + 1 , r t + 1 | s t , a t ) p(s_{t+1},r_{t+1}|s_{t},a_{t}) , as seen in systems like dialogue models, where user responses are inherently unpredictable. On the other hand, deterministic feedback involves no randomness, yielding a fixed next state s t + 1 s_{t+1} and reward r ⁡ ( s t , a t ) r(s_{t},a_{t}) . For instance, when a LLM solves a mathematical problem, the transition is deterministic, where the current state s t s_{t} and action a t a_{t} are combined to produce the next state s t + 1 s_{t+1} .

## 3 Policy Initialization

In reinforcement learning, a policy defines how an agent selects actions in response to environmental states. As discussed in Section 2 , LLMs operate with actions at three granularity levels: solution-level, step-level, and token-level. Solution-level actions represent the coarsest granularity, treating the entire solution as a single action. Step-level actions operate at an intermediate granularity, where individual steps serve as discrete actions. Token-level actions provide the finest granularity, treating each individual token as an action. Taking token-level actions as an example, the action space contains thousands of tokens from the vocabulary, establishing a well-initialized policy becomes essential for effective model performance ( Brown et al., 2020 ) .

As illustrated in Figure 4 , the initialization process of LLMs consists of two primary phases: pre-training and instruction fine-tuning. During pre-training, models develop fundamental language understanding through self-supervised learning on large-scale web corpora ( Sun et al., 2024d ; Weber et al., 2024 ; Liu et al., 2024f ) , following established power-law relationships between computational resources and performance ( Kaplan et al., 2020 ; Hoffmann et al., 2022 ) . Instruction fine-tuning then transforms LLMs from simple next-token prediction to generating human-aligned responses ( Wei et al., 2022a ; Chung et al., 2024 ) . For models like o1, incorporating human-like reasoning behaviors is crucial to enable more sophisticated exploration of solution spaces. We summarize six key behaviors that can be activated through prompts or learned through expert trajectories distillation from LLMs.

### 3.1 Pre-Training

Pre-training establishes basic language understanding and reasoning capabilities in LLMs through exposure to massive text corpora ( Radford & Narasimhan, 2018 ; Lee et al., 2024 ) . For o1-like models, these core competencies serve as the basis for advanced behaviors developed through subsequent learning and search.

#### 3.1.1 Language Understanding and Generation

Pre-training cultivates diverse language capabilities through extensive exposure to natural language ( Radford & Narasimhan, 2018 ) . At the syntactic level, models learn grammatical structures ranging from basic word order patterns to complex dependencies ( Manning, 2022 ) . This syntactic foundation enables pragmatic understanding, including discourse markers and contextual language use, allowing models to adapt different styles across tasks ( Dam et al., 2024 ) . Generation capabilities progress from basic grammatical coherence to sophisticated features like long-range consistency and complex narrative structures ( Tian et al., 2024b ) . Through multilingual training data, models develop cross-lingual abilities, enabling zero-shot transfer and cultural understanding across languages ( Scao et al., 2022 ; Alves et al., 2024 ) . Research shows language understanding emerges hierarchically: syntactic patterns appear early, while logical consistency and abstract reasoning develop later, suggesting the importance of training duration and data composition beyond model scaling ( He et al., 2024 ; Ahuja et al., 2024 ) .

#### 3.1.2 World Knowledge Acquisition and Storage

Pre-training enables comprehensive knowledge acquisition across factual, procedural, and conceptual domains through diverse corpora processing ( Radford et al., 2019 ; Brown et al., 2020 ) . Models develop rich semantic networks of factual knowledge from encyclopedic sources and academic literature, enabling cross-domain reasoning and novel insights ( Chang et al., 2024 ) . Domain expertise emerges from specialized technical content, manifesting in advanced capabilities like mathematical proofs and scientific analysis ( Shao et al., 2024 ; Yang et al., 2024 ) . Procedural knowledge develops through exposure to instructional content and programming languages, enhancing systematic problem-solving abilities ( Ruis et al., 2024 ; Brahman et al., 2023 ; Haluptzok et al., 2023 ) . Mathematical and logical foundations form through formal mathematical texts, establishing logical reasoning capabilities ( Sun et al., 2023 ; Ahn et al., 2024 ) . Recent studies demonstrate that knowledge storage exhibits efficient compression and generalization properties ( An et al., 2024 ) , with abstract concepts requiring more extensive training compared to factual knowledge ( Allen-Zhu & Li, 2024 ) .

#### 3.1.3 Basic Reasoning Abilities