Scaling Laws for Pre-training Agents and World Models
Authors: Tim Pearce
Tabish Rashid
David Bignell
Raluca Georgescu
Sam Devlin
Katja Hofmann
Venue: ICML 2025 poster (2025)
Keywords: world modeling
imitation learning
scaling laws
Primary area: deep_learning->generative_models_and_autoencoders

Abstract: The performance of embodied agents has been shown to improve by increasing model parameters, dataset size, and compute. This has been demonstrated in domains from robotics to video games, when generative learning objectives on offline datasets (pre-training) are used to model an agent's behavior (imitation learning) or their environment (world modeling). This paper characterizes the role of scale in these tasks more precisely. Going beyond the simple intuition that `bigger is better', we show that the same types of power laws found in language modeling also arise in world modeling and imitation learning (e.g. between loss and optimal model size). However, the coefficients of these laws are heavily influenced by the tokenizer, task \& architecture -- this has important implications on the optimal sizing of models and data.

Scaling Laws for Pre-training Agents and World Models
Tim Pearce" Tabish Rashid David Bignell'
Raluca Georgescu' Sam Devlin' Katja Hofmann'
Abstract
In comparison, less is understood about scaling in embod- ied AI. Recent high-impact works show increasing model and dataset size can lead to ever more capable agents for two pre-training objectives; behavior cloning (BC) (Reed et al., 2022; Baker et al., 2022; Brohan et al., 2023) and world modeling (WM) (Hafner et al., 2020; Hu et al., 2023; Yang et al., 2023; Bruce et al., 2024). Such works typically demonstrate the benefit of scale through ablations over few model sizes, shown in terms of downstream agent perfor- mance, confirming the intuition that 'bigger is better' (Sartor & Thompson (2024) provide an aggregated analysis). How- ever, this leaves a large gap to the precise understanding of scale in LLMs, where for a given increase in compute, mod- els can be sized optimally, and their performance accurately predicted.
The performance of embodied agents has been shown to improve by increasing model parame- ters, dataset size, and compute. This has been demonstrated in domains from robotics to video games, when generative learning objectives on offline datasets (pre-training) are used to model an agent's 's behavior (imitation learning) or their environment (world modeling). This paper char- acterizes the role of scale in these tasks more precisely. Going beyond the simple intuition that 'bigger better', we show that the same types of power laws found in language modeling also arise in world modeling and imitation learning (e.g. be- tween loss and optimal model size). However, the coefficients of these laws are heavily influenced by the tokenizer, task & architecture- this has important implications on the optimal sizing of models and data.
This paper helps close this gap. Similar the study of scale in LLMs, we focus on the effect of scaling on generative pre-training loss (rather than on downstream agent perfor- mance, or reward- or representation the infinite data regime, on a fixed offline dataset. Under this setting, we train families of transformers on next-token prediction tasks using architectures popular in both world modeling and BC tasks. This leads several contributions, summarized in Figure 1.
objectives), in
1. Introduction
Much progress in Al in the early 2020's has been driven by increasing model size, dataset size, and training com- pute. Whilst conceptually simple, the importance of this practice has led to an emerging subfield studying the sci- ence of scaling. This field answers questions such as how to estimate the benefit of increased compute investment, or how to optimally trade-off model and dataset size.
For the first time, we show scaling laws similar to those in LLMs can be observed in world modeling with tokenized observations and actions (Section 4.1, Figure 1a).
The optimal trade-off between model and dataset size in world modeling is influenced by the tokenizer's 's compres- sion rate (number of tokens per observation) (Section 4.1, Figure la & b).
The role of scale in pre-training is so far best understood in large language models (LLMs). Following the observation that the empirical relationship between loss and key seal- ing quantities can be accurately described by power laws (Kaplan et al., 2020), works studied the precise trade-off between model and dataset size (Hoffmann et al., 2022), as well as considering inference compute (Sardana & Frankle, 2023), repeated epochs (Muennighoff al., 2024), parame- ter counting (Pearce & Song, 2024) etc. (Section 2.1).
Scaling laws for BC with tokenized observations are harder to observe under modest compute budgets. The optimal trade-off favors smaller models and more data (Section 4.2, Figure Ic).
Scaling laws similar to those in LLMs can once again be observed in BC with one continuous encoding per observation (Section 4.2, Figure Id).
'Equal contribution 'Microsoft Research (SD now at Meta). Correspondence to: Tim Pearce <.>.
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
Our findings can be understood through small-scale lan- guage modeling experiments (Section 5).

---

Scaling Laws for Pre-training Agents and World Models
(a) WM-Token
(b) WM-Tookn 550
5 10
N Nopimal
X
C
Noptinal O CO.62
5 x 10
D Doptiaal C
D Doptimul C0
x 10
Loss
Loss
4 10
15M param
1M param
27M parame
2TM peame
x 10
52M parans
S2M pazma
IECM porns
180M panms
2063 parmss
206M parans
Efficiat Sounier
Effiicas Bottier
L-0.1 C7
L-0.1 C
10"
10"
10"
10
3 10
10"
10
FLOPs
FLOPs
(c) BC-TTkee-5444
(d) BC-CNN
4.8 10
Nopumaa ok C
4 x 10
Noguimal x Co.ss
D Dogeimal
0
3.9 10
D Dopimal X C0.34
4.6 16x
3.8 10
Loss
Loss
4.4x
3.7 x 10
1M porms
2M poranc porass
10M panzs
4.2x
IIM parenc
3.6 10
SIM panes
20M ponns
Edficient Gontier
3.5 x 10
fountirr
0.1
L- 3.56 - -
4x 10
10"
10"
10" FLOPs
10"
10
10"
10
10"
10
FLOPs
Figure 1. This paper observes that scaling laws, as originally found in LLMs, also emerge in the tasks of world modeling and BC, when studying the pre-training loss on large datasets of human behavior. (a. b) For world modeling, the power law coefficient determining optimal model size is affected by the compression rate of the tokenizer. (c) In BC with tokenized image observations (BC-Token), small models need a large FLOPs budget to saturate, making these scaling laws less clear cut. (d) Moving to single continuous embedding per observation remedies this (BC-CNN), producing prototypical scaling laws and a more balanced optimal model size coefficient. Note BC-Token & BC-CNN losses are not directly comparable.
Organization. Section 2.1 provides detailed related work, contrasting the current understanding of scaling in embodied AI with other domains, and justifying pre- pre-training loss as proxy for online reward. Section 3 introduces details for our main experiments, including the architectures & datasets considered, and details of scaling laws analyses. Section 4 presents our main results in world modeling and BC. Section 5 presents insights behind our main results, including a set of scale language experiments mimicking aspects of our main experiments. Section 6 discusses our findings and notes limitations.
2.1. Related Work
Scaling laws origin. The term scaling laws is used through- out the engincering and physical sciences to denote power law relationships between two quantities, e.g. duration of a voleanic eruption and the probability of continuing (Can- navà & Nunnari, 2016). The name derives from the scale- invariant property of power laws. While early work sug- gested that power laws could be good empirical descriptors of pre-training loss in deep learning (Hestness et al., 2017: 7; Rosenfeld et al., 2019), Kaplan et al. (2020) provided a comprehensive study of power laws in transformer LLMs, and popularized the usage of scaling laws in this context.
2. Background
Scaling laws in LLMs. As the real-world value of LLMs was understood, scaling in LLMs became a high-priority research topic. Hoffmann et al. (2022) conducted a pre-
This section introduces related work, and outlines arguments and evidence supporting using pre-training loss to study scaling in embodied AI.
For two variables x & y, the power law y ax is invariant to scaling by constant c. Formally: a(cz)) d' d'rr" = c'y = daa = y = ax
2

---

Scaling Laws for Pre-training Agents and World Models
cise analysis into the trade-off of model and dataset size, finding they should be increased in equal proportions. This conflicted with the suggestion that model size should be pri- oritized (Kaplan et al., 2020) an incorrect conclusion that Pearce & Song (2024) showed largely arose from counting only non-embedding parameters.
the benefit of scale is generally shown through increasing model size on a specific downstream task of interest (e.g. measured by completion rate)- aggregated survey is provided by (Sartor & Thompson, 2024).
Tuyls et al. (2023) offer a valuable initial investigation into scaling laws for BC. They fit power laws to both BC pre- training loss and online return, when scaling width of single- layer LSTM models on datasets generated by fixed high- reward policies. We extend this line of investigation by studying transformer models trained on datasets of human behavior, discovering effects of architecture choices on scal- ing coefficients. In addition we study scaling laws in world models for the first time.
Many other aspects of LLM scaling analyses are beginning to be refined. Su al. (2024) revisited the methodology for finding scaling coefficients. Hägele et al. (2024) found that multiple independent cosine schedules could be reproduced more efficiently through a constant learning rate with mul- tiple short decays, or stochastic weight averaging. Pearce & Song (2024) & Porian et al. (2024) found that well-tuned constant learning rates were sufficient to recover certain coefficients. Bi et al. (2024) study the effect of various hy- perparameters on scaling. Muennighoff et al. (2024) looked at repeated epochs, finding up to four epochs produce negli- gible departures from the infinite data regime. Sardana & Frankle (2023) factored in inference compute to the defini- tion of compute-optimal. Isik et al. (2024) study the link between pre-training loss and downstream performance. A further line of research aims to explain why power laws are such a good descriptor of empirical deep learning (Hutter, 2021; Maloney et al., 2022; Bahri et al., 2024).
2.2. Pre-training Loss as a Proxy for Performance
A major difference between scaling research in LLMs and embodied AI that LLM research uses pre-training loss as the main variable of interest, while embodied Al has focused on downstream online task performance. This handicaps em- bodied AI scaling research measuring online performance for a single model checkpoint in complex environments like robotics or modern video games is expensive in time and hardware, requiring multiple repeated runs to allow statis- tically significant comparisons. Furthermore, models may first require a period of fine-tuning before evaluation. By contrast, pre-training loss is available for free at any point of a model's training.
Scaling laws in image and video generation. Scaling laws have also been observed in auto-regressive modeling of video and images (Henighan et al., 2020; Tian et al., 2024). Henighan et al. (2020) found the optimal trade off between model and dataset size to match their reported LLM coef- ficient (Neptimul x CO.7) and was not affected by tokenizer. Our experiments offer different findings in the domain of world modeling - using updated methodologies to measure this trade-off, we find it is affected by the tokenizer.
We believe embodied Al's focus stems from reports that validation loss is only weakly correlated with online per- formance (Hussenot et al., 2021; Li et al., 2024). However, such observations have been made with fixed-sized training datasets and held out validation sets, where effects of over- fitting may be slightly beneficial. In contrast, scaling law studies are conducted in infinite data regime, where dat- apoints are not trained more than once, making train and test losses equivalent, and overfitting effects not applicable.
Scaling in embodied AI. Compared to LLMs, an under- standing of scale in embodied settings is less advanced. Early successes in competitive games showed that reinforce- ment learning (RL) at scale could surpass human perfor- mance (Silver et al., 2017; Berner et al., 2019). In self-play RL, power laws were observed between certain quantities by (Neumann & Gros, 2022). Meanwhile, Hilton et al. (2023) noted that reward signals do not generally follow power laws, and defined transformation of reward (intrinsic per- formance) that create self-consistent scaling laws.
To evidence that pre-training loss can be a good proxy for online return in the infinite data regime, we conducted a anta-anaaysis of Tuyls al. (2023), who were able to roll out large number of checkpoints for two reasons. 1) They used simple lightweight environments (Atari & NetHack). 2) Their demonstration policy was high-skill, removing any need for fine-tuning. Figure 2 plots online environment return vs. pre-training loss for several environments (com- puted by tabulating pairs of points from Figure & 10 in Tuyls et al. (2023)). The correlation coefficient for all envi- ronments is stronger than -0.94. Figure 3 shares evidence from experiments that pre-training loss is well corre- lated with the video-generation quality of world models- providing correlation coefficients around 0.8. Further details in Appendix D.
Inspired by the effectiveness of scaling in LLMs, embodied Al research has recently begun to explore the effective- ness of generative pre-training objectives on offline datasets, when executed at scale. This includes behavior cloning objectives in video games (Baker et al., 2022; Raad et al., 2024), robotics (Brohan et al., 2022: 2023; Padalkar al., 2023; Bousmalis et al., 2023), or multiple domains (Reed et al., 2022), as well as world modeling objectives (Hu et al., 2023; Yang et al., 2023; Bruce et al., 2024). In these studies,
3

---

Scaling Laws for Pre-training Agents and World Models
NetHack
Bank Heist
Breakout
Space Invaders
10
Corelaien R-4837
Comilason, Re-d
4x10
Comiiiien
3x10
10"
Online retumm
log pfride sizs)
Online petum
lngwomedde stz)
Online retum
log in moniel sune)
t0
Online
2 x110
6
10
4 10"
10
10'
10 Pre-training loss
10
10
10
10" Pre-training loss
Pro-training loss
Pre-training loss
Figure 2. Our meta-anaa anallsis of Tuyls et al. (2023) evidences that pre- pre-training loss is strongly correlated with reward in BC tasks when in the infinite data regime.
More intuitively, improving a next-token prediction loss in BC and WM requires models to 'know more' about behav- iors and the environment, creating more useful pre-trained checkpoints for specialization to downstream tasks. In BC, better predicting the next action dataset of human be- havior requires understanding the objectives humans are trying to complete, alternative social behaviors they might choose perform, as well making in-context inferences about the skill level and mental state of individuals. In WM, decreasing next-token prediction loss might follow cur- riculum, first requiring a model to capture basic shapes and colors, then textures and physics, followed by rare object interactions, and finally even complex stochastic elements the environment such other intelligent agents.
Given this data format, we consider two tasks. World model- ing (WM) (Ha & Schmidhuber, 2018) predicts future obser- vations from previous observations and actions. This allows an agent to explicitly understand how its environment works, which can be used for planning, or dyna-style RL (Sutton, 2018). 8). Behavior cloning (BC) predicts the future actions that the dataset's 's demonstrators take (Bakker et al., 1996). This creates a policy that can be directly used environment, either as-is or following further fine-tuning. Concretely, these two tasks require modeling the following quantities,
in the
World modeling: P(o P(o ..... Or-k; a.. at-k); (1)
Behavior cloning: P(a:lo.... ... Or-k, a1-1 ... at-k-1
(2)
3. Methodology
This work focuses on generative pre-training aiming to model this full conditional probability distribution. Weleave a study of scaling laws for alternative objectives, e.g., explic- itly targeting representation learning (Nair et al., 2022) or reward-centrie models (Hafner et al., 2020), to future work.
This section provides details for our main experiments. We describe the pre-training tasks, architectures, and datasets considered. We also detail the methodology used in the scaling law analyses.
3.2. Architectures
3.1. Tasks
All experiments revolve around GPT-2 style causal trans- formers (Radford al., 2019) as the core of the model. However consider two different methods for inputting image observations, summarized in Figure 4. Section 3.4 details how we measure the model size of each.
We consider trajectories constructed as sequences of alter- nating observations o, and actions a for timestep t N. In this work, observations are always images, O, E 33xwxx and any continuous actions discretized during prepro- cessing leaving, a& E (0, 1)
Tokenized architecture. Our first architecture tokenizes each image observation into multiple discrete tokens. This is done with a frozen VQGAN encoder Encę(o;) + Z,, where z, {1, 2, Vo) for vocabulary size V, and la- tent dimension d.. Discretized actions are mapped non-overlapping vocabulary. Following tokenization, train- ing sequences take the form,
10'
Sizz I5M
Sizee 52M
Sizz 134
10
Sixc MMIM
[=] -... =( de ob...., ... de
(3)
Correlation,
0.83
241
10
3.55
2.50
2.55
3.0
Loos
325
3.50
3.75
4.00
where each item of the sequence an integer within our vocabulary. A transformer is then trained to maximize the
Figure 3. Our experiments suggest pre-training loss is a good proxy for generation quality of the world model, here FVD.

---

Scaling Laws for Pre-training Agents and World Models
W WM-Toeen
BC-Token
BC-CNN
X\/
Transformer
Transformer
Transformer
2,
Zi+1
a1+1
KUBAA
VOGGNN
Oz+1
OH+
O2 Figure 4. The World Modelling (WM) and Behavior Cloning (BC) tasks architecture combinations considered in this work. The fire symbol signifies trainable components, the ice symbol signifies frozen pre-trained components.
Oct1
Oc11
likelihood of either the latent image tokens (world model- ing), action tokens (BC).
paced 4 vs 4 multiplayer game, with a range of characters, abilities and maps. Game play is highly complex due to cooperative and competitive dynamics. Success requires good high-level strategies (e.g. choosing which map regions to fight for), as well as fine-grained reactive control during combat. Figure 11 shows example dataset sequences.
This tokenized architecture is widely used both in world modeling (Micheli et al., 2022) and BC tasks (Bousmalis et al., 2023). Gato (Reed et al., 2022) used a similar design but with continuous patches rather than discrete tokens. Our implementation tests both a 'small' (28M parameters, d= = 256) and 'large' (150M parameters, d_ 540) VQGAN further details in Appendix A.
Supported by the game's developer Ninja Theory, we com- piled dataset of over 7 years of anonymized game play, containing both image observations and controller actions. We refer to this as the 7 map dataset. We also use subset of this for some experiments, of around one year from single map, which we name the Sky Garden dataset. Appendix A.3 provides further details.
CNN architecture. Our second architecture differs in two ways. 1) Each image observation is input into the trans- former as a single continuous embedding, extracted from a small trainable convolutional neural network (CNN). 2) Action dimensions are predicted independently (rather than in series), assuming ..a) Rs II P(afl...) One transformer forward pass produces an action prediction.
As a secondary dataset we use RT-1 (Brohan et al., 2022), comprising 14 days of humans operating robotic arm range of manipulation tasks such as "pick banana from white bowl'. Using this smaller dataset allows us both to verify that conclusions on the large-scale vidco games dataset hold in a real-world robotics domain, and also allows us to run several small scale ablations in WM. Appendix provides further details about the dataset and tokenizers used.
This produces architecture similar to (Baker et al., 2022) (VPT additionally used transformer-XL. and refined hier- archical action space). Our implementation uses an Impala- style (Espeholt et al., 2018) CNN with 0.6M parameters for embedding image observations.
3.4. Scaling Analysis Methodology
3.3. Datasets
We study the relationship between several quantities.
This paper focuses on the effect of scaling on the pre- training loss over an offline dataset. To study this cleanly, datasets must meet two criteria.
Model size N. the total number of trainable param- eters (ignoring VQGAN parameters for WM-Token & BC-Token, but including the fixed-size CNN for BC-CNN), Embedding parameters are included in the count following (Pearce & Song, 2024).
1. Dataset size. Repeated training on the same data alters the effect of scaling datasets should be large enough that training is done in the infinite data regime.
Dataset size D, the total number of inputs the trans- former sees during training. For WM-Token and BC- Token this is d, + de per observation & action pair, for BC-CNN this is one per observation & action pair.
2. Dataset diversity. The behavior and environment must contain enough richness and variety that pre-training loss does not saturate across model sizes tested.
Many existing benchmark datasets fail to fulfill these criteria either due to limited size, or because behavior is generated from a fixed policy, or the environment is too simple.
Compute C, the number of floating point operations (FLOPs) used during training. The common approxi- mation of C GN (Kaplan et al., 2020) is used.
Our work primarily focuses on a dataset of human behavior collected in a video game Bleeding Edge. This is a fast-
Loss L, the standard classification cross- cross-entropy loss (all targets are discretized). Note training loss is an
5

---

Scaling Laws for Pre-training Agents and World Models
accurate proxy for test loss when in the infinite data regime (Appendix 1.3.1 analyzes further).
Table 1. Summary of fitted scaling coefficients for our main experi- ments. Note that we favor the Frontier fit when available, and only use the Parametric fit for BC- Token-54p (see Section 3.4).
More specifically, we are interested in 'compute-optima" versions of each quantity. For loss, this defined as the minimal loss possible for given FLOPs budget.
Frontier fit
Parametric fit
Experiment
Neptiaa DC C Doptuas C
Nopties X C Deptrss C
WM-Tokeen2256
0.49
0.52
0.48
Loptmal (C)
min D), s.t. C=6ND
(4)
PMM-Token-540
0.62
0.37
0.78
0.22
BC-Token-540
0.32
BC-CNN
0.66
0.34
0.47
0.53
where L(N,D) D) is the empirical loss achieved with an N parameter model trained D tokens. We further define optimal model and dataset sizes as the configuration that produce this minimal loss given a FLOPs budget,
We again use the curve_ fit function, fitted to models along the efficient frontier. During fitting, we set bounds on the variables: co E [0, o0], [-1, 1], E oo].
NNopiiall (C), Doptimal (C) (C) N. st. CN6ND
argmin D), (5)
Training details. While early scaling studies conducted sweeps over multiple cosine decays of differing lengths (Kaplan et al., 2020; Hoffmman et al., 2022), follow up work found this redundant (Pearce & Song, 2024; Hägele et al., 2024; Porian et al., 2024). We follow the approach of using constant learning rate per model, so each requires only one training run. We aim train models until they have passed their compute efficient FLOPs budget. We only modify the parameters of the transformer, following the configurations documented in Appendix A.
Scaling analysis. Scaling law analyses fit power law rela- tionships predicting these compute-optimal quantities. For predicting optimal model and dataset size, we use.
Noptimal (C) = Ca Doptimal (C) = bn,,,,
with fitted constants ao; a, bo, We consider two methods to fit these relationships, introduced by (Hoffmann et al., 2022). Their Method I, which we term Frontier fit, classi- fies efficient models as those falling on the efficient frontier (see Figure 1). Coefficients can then be estimated straight- forwardly through a line of best fit on a plot of FLOPs vs parameters or data for these efficient models.
4. Scaling Analysis in Embodied AI
This section presents our main results. We begin by con- sidering the scaling laws for the task of world modelling in Section 4.1 with two different tokenizers (turning image observations into 256 and 540 tokens for the small and large variants respectively). Section 4.2 then considers the task of BC both with tokenized and CNN architectures. Finally, Section 4.3 tests the extrapolation capability of these scaling laws for the task of world modeling.
Frontier fit our preferred method when available it avoids making any assumptions about the training curves, directly fitting the best models observed. However, it re- quires training models past the point where they are the optimal configuration (seen on loss-FLOPs plot as over- lapping curves). In some of our experiments (BC-Token and Section 5.1), this was not possible.
In these situations, we resort to Method 3 of (Hoffmann et al., 2022), which term Parametric fit. This fits the coefficients o; B, Ne, De, E to parametric loss form,
4.1. Scaling Analysis in World Modeling
Figures 5 & 6 present our results for the task of world modeling, with the scaling law coefficients summarized in Table 1. For WM-Token-256 coefficients for model and dataset size are both 0.5, e.g. one should increase both model and dataset size in the same proportions. This matches the scaling laws observed in L.MSs (Hoffmann et al., 2022). Increasing the number of tokens per image to 540 for WM-Token-5544 changes the optimal trade-off between model and dataset size, skewing towards model size; Noptimal = 0.62, Doptimal = 0.37. We discuss this further in Section 5.3.
No De L(N. D) = Na D E,
+
find that the optimal
to the empirical training curves. In our implementation, we use SciPy's curve_fit function. We find 3/(a + b) = a/(a + 3). This makes a very strong assumption about the training curves, but allows coefficients to be estimated at a smaller compute budget.
For loss prediction we use the form recommended by Pearce & Song (2024).
Enppimal NN, D) coC C + E.
(8)
Appendix Figure 15 visualizes the power law fits for the RT- robotics dataset, confirming that this predictable scaling behavior is not specific to human behavior in video games, and also emerges on world robotics tasks with high-skill human operators.
Note that by subscribing to C' = 6N D we find a = - b:
N x Ca = C/D x C = D x Cl-a. Hence, at
times we only deseribe relationships in terms of N x C", with N x D³ l- implied.
6

---

Scaling Laws for Pre-training Agents and World Models
6-10
770
Fane Las vi.i
les Incal
Oytinal tukena
3ex10
Optinaa
3x10
10
FLOPs
Figure 5. WM-Token scaling with da 256 tokens per image observation. Left shows the parametric fit. Middle & right show the frontier estimating optimal model & dataset size respectively.
6-10
10
270
5x10
1100
Optimal tokess
Loa
10
4x10
10
10"
10
10
FLOPS
FLON
FLOR.
Figure 6. WM-Token scaling with dg =540 tokens per image observation. Layout as in Figure 5. Compared the results for WM-Token- 256, the power law coefficient for Naptial increases from 0.49 to 0.62.
4.2. Scaling Analysis in Behavior Cloning
4.3. Extrapolation in World Modeling
We present our results on the scaling law coefficients for BC-Token in Figure 7. Despite sharing architecture with WM-Token-550 we observe the opposite dependence on model and dataset sizes. The coefficients skew heavily towards dataset size; Noptimal = 0.32, D Doptimal = 0.68 (com-
To test the extrapolation accuracy of our derived scaling laws, train a 894M parameter WM-Token-225 model with order of magnitude more compute than used for the scaling law analyses. Figure 9 presents both the learning curve as well as the extrapolated lines derived from the Fron- tier method. We take the point with the loss value clos- est to our extrapolated loss curve (~ 1.58 x 102! FL.OS), and mark it on the Frontier fit extrapolations. We observe very good agreement between that point and our compute- optimal predictions for both model and dataset size, demon- strating the accuracy of our derived scaling laws. The gap between our prediction and the actual training run suggests we could further optimize the hyperparameters (learning rate and batch size in particular) for the 894M parameter model, which was not extensively tuned due to compute requirements.
pared to NNotimall = 0.62, Doptiall = 0.37 explained in Section 5.1). Furthermore, under the same compute budget the compute-optimal model sizes are significantly smaller. For compute budget of 10 and 1019 FLOPs we find that model sizes of 2M and 11M are compute-optimal for BC- Token-540 compared to 27M and 110M for WVM-Token550 In our experiments, we observe the losses for the BC-Token models take much longer to plateau leading to less overlap between model sizes. Hence, the frontier fit is not suitable for accurately estimating the scaling law coefficients, and rely the parametric fir for these results. Appendix E investigated whether smaller models would allow estimation via frontier fit, but this was unsuccessful.
5. Further Analysis
To better understand the change in the scaling law coeffi- cients, we now consider the BC-CNN architecture for the task of BC in Figure 8. For this architecture, we observe that the coefficients now skew towards model size (similarly to those in (Tuyls et al., 2023)), with Nopimal 0.66, and Doptimal 0.34. Section 5.2 provides more intuition the differences between the WM-TToke and BC-Token setups that lead to this change.
Section 4 made several observations about the effect of scale in the pre-training of embodied agents. This section aims to understand these results further, and provide intuition for why they occur. Specifically we target three questions.
Q1: Why does BC-Token produce training curves that do not plateau, while WM-Token does, given an identi- cal architecture and dataset? (Section 5.1)
Further to studying the differences in scaling law coeffi- cients between tasks and architectures, we also study the accuracy of extrapolation.
Q2: Why does moving from BC-Token to BC-CNN resolve this issue? (Section 5.2)
Q3: Why does increasing the tokens per image obser- vation (256 540) lead to an increase the optimal model size coefficient (0.49 to 0.62)? (Section 5.3)
7

---

Scaling Laws for Pre-training Agents and World Models
16
Oytinal tukena
5.10
Optinaa
16
4x10
FLOPs
Figure 7. BC-Token scaling with d2 540 tokens per image observation. Layout as in Figure 5. Models above 2M parameters do not saturate over the FLOPs range considered and coefficients not be inferred using the frontier fit method.
42 16
230
10
Optanull tokzus
16'
18
10
10
10"
10"
FLOP
FLONs
FL.0%
Figure 8. BC-CNN scaling. Layout Figure 5. Compared to the results for BC-Token the model sizes considered compute-optimal
are considerably larger. The power law coefficient for Noptimal also increases from 0.32 to 0.66 skewing towards scaling model size as opposed to dataset size when scaling up compute.
10
Fanr aid
10a
Toeer Inc.))
4<10f
Optinal tolmm
3x10
10'
10'
1.0
10
10
10"
10
10"
FLON
FLON
FLOP
Figure 9. Testing extrapolation capability of our derived scaling law for WM-Token-2556 by training an 894M parameter model with an order of magnitude more compute than was used for the scaling law analyses. We observe good agreement between our predicted optimal loss/model size/number training tokens (dotted lines) and our actual training run marked by 'x'.
Dense loss
Sparse loss
Sparse loss, super-classed Naptiall x 1. Dopimal 85
N Nopimall oc 00.63 D Deptimal 0..37
Naptiall
0. D Doptimal x 5
3x
1-
5180 poum
12
2
5x
4-10
10"
10
10
FLOPs
FLOPs
FLOPS
Figure 10. Training curves and parametric fit for character modeling experiments. The standard dense LLM loss has been modified to reflect properties of BC sparse loss (one of tokens), and then additionally super-classing the targets into two classes.
5.1. Q1: BC-Token vs. WM-Token
color and texture for an image patch. Many vocabulary items may only be used to model specific map regions or special abilities. Hence, the world modeling loss is very granular. On the other hand, player take the same action in multiple different situations continue straight could be used to escape enemy, chase an enemy, or navigate to a checkpoint. Hence, the supervision for BC is more vague and abstracted. We can think of this as a super-classed label.
The lack of saturation of BC-Token models compared to WM-Token models can be attributed to two factors. The first sparser loss. A single observation-an action pair is discretized into +d, total tokens. With the large VQGAN tokenizer, world modeling receives supervision for ds/ + d.)) = 540/556 x 97% tokens, while BC is supervised for dd,(d, + do) = 16/556 x 3% of tokens.
The second factor is the granularity of the targets. The large tokenizer creates a world modeling vocabulary size 4096. Each vocabulary item roughly corresponds to a specific
To demonstrate the effect of these two factors on optimal model size coefficients, we run a set of tiny-scale experi-
8

---

Scaling Laws for Pre-training Agents and World Models
ments in language modeling. Transformers are trained on next-character prediction, dataset of Shakespeare text² using single character for each token. Model sizes are var- ied from 4k parameters to 17M parameters. Context length is fixed at 16 characterstookens.
of examples would need be seen (smaller dataset size). It is less clear what ratios these ingredients decrease in, and hence what effect a lower compression rate has on the optimal model size coefficient.
Using the small scale RT-1 dataset, we conduct a more thorough investigation of the effect of tokens-per-image observation on scaling coefficients. First we train a range of image tokenizers with 2o E [16, [16,36, 64, 100, 256). vi- sualized in Appendix Figure 14. For each VQVAE, we then train a range of WM-Token model sizes N E [0.08.M, 0.2MM, 0.28M, 0.54M, 0.99MM], and measure scal- ing coefficients using the frontier fit method, repeating three times.
Figure 10 (left) shows training curves over all 16 tokens, followed by a sparse loss where supervision is only provided from the final token (middle), and then additionally under super-classed setting (right). This super-classes the final target- rather than using all 128 ASCII characters, they are randomly shuffled into one of two macro classes.
These modifications are intended to mirror the effect of moving from WM-Token to BC-Tokkn. We compute op- timal model size coefficients using the parametric fit as most models are not trained long enough for the frontier fit method. Indeed, we see that the coefficient drops from 0.63 to 0.15 with both the sparse and super-classed loss. This matches the magnitude of decrease seen in Table 1 from 0.66 (BC-CNN frontier fit) to 0.32 (BC-Token- paramet- ric fit), indicating that the proposed mechanisms explain our findings.
Figure 16 plots all coefficient vs. tokens-pen seren--prr-iaae-- we observe that the optimal parameter scaling coefficient in- creases with decreasing compression.
To investigate whether compression affects the optimal model size coefficient outside of embodied domains, we ran a small scale experiment in language modeling using two text representations; 1) ASCII character-level tokeniza- tion. (low compression) 2) GPT-2 tokenizer (high compres- sion). We used the BookCorpus dataset (Zhu al., 2015), and trained models past their compute-optimal point, the Frontier fit method could be used for coefficient estimation.
5.2. Q2: BC-Token vs. BC-CNN
Despite the same non-granular loss signal, why does switch- ing architecture from BC-Token to BC-CNN makes the loss of similar model sizes plateau under much smaller compute budget?
Appendix B shows results. Under the character-level tok- enizer (low compression), we find Noptimal xx C0.66 For the GPT-2 tokenizer (high compression), we find Naptimal X C Hence, in language the more compressed representa- tion also leads lower optimal model size coefficient.
Consider each architecture using a transformer with 1M parameters. Observe from Figure 4 that BC-Token receives d₂ + dg = 556 inputs for every action å, it predicts, while BC-CNN receives just one input for every action predicted. Hence, BC-Token uses around 556 times more compute in x 10 FLOPs) than BC-CNN (1 x 2 1M s 2 x 10t FLOPs). This means that even with the same number of parameters, BC-Token can learn a far more expressive function than BC-CNN. Hence, BC-Token requires far more tokens to match this expressivity, and training curves for a given model size plateau much later.
6. Discussion & Conclusion
its action prediction (556 2 1M
This paper establishes a deeper understanding of scaling laws for world modeling and behavior cloning, two tasks that underpin embodied AI applications in domains such as video games and robotics. Focusing on generative pre- training of such models, we show that it is possible to re- cover scaling laws similar to those established for LLMs. Establishing this link is key to making efficient use of avail- able resources in training compute-optimal models.
5.3. Q3: WM-Token-256 vs. WM-Token-550
In world modeling, we show models can be smoothly sealed following best practices and insights from the LLM liter- ature. We further establish that scaling coefficients are affected by the tokenizer's 's compression rate. In BC, the choice of architecture greatly impacts scaling coefficients. For architectures with tokenized image observations, dataset size should be increased much more rapidly than model size. Mcanwhile, for BC-CNN architectures, model size should be increased faster than dataset size.
Finally, we seek to understand why the optimal model size coefficient increases when moving from the 256 the 540 token VQGAN. As the number of tokens per image obser- vation are increased, the compression rate of the tokenized representation decreases. We would expect that each individ- ual token becomes easier to predict in this less compressed representation. This would mean a less expressive function is needed (smaller model size). but also a smaller number
Shakesseaaee character dataset from: https:/ //tgth/b..t com/ katpatta/aaaaaoocP
9

---

Scaling Laws for Pre-training Agents and World Models
Impact statement
al. Robocat: A self-improving foundation agent for robotic manipulation. arXiv preprint arXiv:2306. 1706, 2023.
Limitations. While we show that scaling laws emerge in the infinite data regime for generative pre-training ob- jectives in certain embodied datasets, future work could investigate scaling laws alternative objectives and further datatsets. We have provided early evidence that this pre- training loss strongly correlates with online performance, but future work could seek further evidence, particularly in real-world robotic experiments.
Brohan, A., Brown, N., Carbajal, J., Chebotar, Y., Dabis, J.,
Finn, C., Gopalakrishnan, K., Hausman, K., Herzog, A.,
Hsu, J., et al. Rt-1: Robotics transformer for real-world control at scale. arXiv preprint arXiv:2212.068/ 2022.
Brohan, A., Brown, N., Carbajal, J., Chebotar, Y., Chen, X., Choromanski, K., Ding, T., Driess, D., Dubey, A., Finn, C., et al. Rt-2: Vision-language-ace
Our observation that scaling laws arise in embodied Al suggests great potential in the direct scaling up of existing techniques. In consequence, it suggests shift in focus towards data collection - how much data need be collected and with what strategies? Finally, we note our analysis has not factored in additional real-world constraints such as inference time.
models transfer web knowledge to robotic control. arXiv preprint arXiv:2307.1588 2023.
Bruce, J., Dennis, M. D., Edwards, A., Parker-Holder J., Shi, Y., Hughes, E., Lai, M., Mavalankar, A., Steiger- wald, R., Apps, C.,, et al. Genie: Generative interactive environments. In Forty-first International Conference on Machine Learning, 2024.
Ethics. Data was provided via a partnership with the Bleed- ing Edge game developer Ninja Theory, who collected large corpus of human gameplay data. Data collection was covered by an End User License Agreement (EULA) and our use of the data was governed by a data sharing agree- ment with the game studio, and approved with an IRB. This data was recorded between September 2020 and October 2022. To minimize risk human subjects, personally iden- tifiable information was removed from the data.
Cannavd, F. and Nunnari, G. On possible unified scaling law for volcanic eruption durations. Scientific reports, (I):22289, 2016.
Espeholt, L.,, Soyer, H., Munos, R., Simonyan, K., Mnih, V., Ward, T., Doron, Y., Firoiu, V., Harley, T., Dunning, I., et al. Impala: Scalable distributed deep-rl with im- portance weighted actor-learner architectures. In Interna- tional conference on machine learning, pp. 1407-1416. PMLR, 2018
References
Bahri, Y., Dyer, E., Kaplan, J., Lee, J., and Sharma, U. Ex- plaining neural scaling laws. Proceedings of the National Academy of Sciences, 121(27):c231 1eeee2211111111 1878121, 2024.
Esser, P., Rombach, R., and Ommer, B. Taming transformers for high-resolution image synthesis. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 1 12873- 2021.
Baker, B., Akkaya, I., Zhokov, P., Huizinga, J., Tang, J., Ecoffet, A., Houghton, B., Sampedro, R., and Clune, J. Video pretraining (vpt): Learning to act by watching unlabeled online videos. Advances in Neural Information Processing Systems, 35:24639-24663 2022.
Falcon, W. and The PoToh Lightning team. PyTorch Light- ning. March 2019. URL /ttpp::/ggttiiie com/ L dighhtnigg-a -g highigggggnnigginnnngggggg iighningg.
Bakker, P., Kuniyoshi, Y., et al. Robot see, robot do: An overview of robot imitation. In AISB96 Workshop on Learning in Robots and Animals, volume 5, pp. 3-11, 1996.
Ha, D. and Schmidhuber, J. World models. arXiv preprint arXiv: 1r03.:/803..0//10 2018
Hafner, D., Lillicrap, T., Norouzi, M., and Ba, J. Mas- tering atari with discrete world models. arXiv preprint arXiv:2010.02//9 2020.
Berner, C., Brockman, G., Chan, B., Cheung, V., Debiak, P., Dennison, C., Farhi, D., Fischer, Q., Hashme, S., Hesse, C., et al. Dota 2 with large scale deep reinforcement learning. arXiv preprint
Hägele, A., Bakouch, E., Kosson, A., Allal, L. B., Von Werra, L., and Jaggi, M. Scaling laws and compute- optimal training beyond fixed training durations. arXiv preprint S4Ki:::200518399. 2024.
2019.
Bi, X., Chen, D., Chen, G., Chen, S., Dai, D., Deng, C.,, Ding, H., Dong, K., Du, Q. Fu, Z., et al. Deepseek Ilm: Scaling open spen-source language models with longtermism. arXiv preprint arXiv:2401.0295.. 2024.
Henighan, T., Kaplan, J., Katz, M., Chen, M., Hesse, C., Jackson, J., Jun, H., Brown, T. B., Dhariwal, P., Gray, S., et al. Scaling laws for autoregressive generative modeling. arXiv preprint arXiv:2010. 14701, 2020.
Bousmalis, K., Vezzani. G., Rao, D., Devin, C., Lee, A. X., Bauza, M., Davchev, T., Zhou, Y., Gupta, A., Raju, A.,
10

---

Scaling Laws for Pre-training Agents and World Models
Hestness, J., Narang, S., Ardalani, N., Diamos, G., Jun, H., Kianinejad, H., Patwary, M. M. A., Yang, Y., and Zhou, Y. Deep learning scaling is predictable, empirically, arXiv preprint arXiv:/712.000
Muennighoff, N., Rush, A., Barak, B., Le Scao, T., Tazi, N., Piktus, A., Pyysalo, S., Wolf, T.,, and Raffel, C. A. Scaling data-constrained language models. Advances in Neural Information Processing Systems, 36, 2024.
2017.
Hilton, J., Tang, J., and Schulman, J. Scaling laws for single-agent reinforcement learning. arXiv preprint arXiv:2301.134444 2023.
Nair, S., Rajeswaran, A., Kumar, V., Finn, C.,, and Gupta, A. R3m: A universal visual representation for robot manipulation. arXiv preprint arXiv:2203, 12601, 2022.
Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T. Rutherford, E.,, Casas, D. d. I.., Hendricks, L. A., Welbl, J., Clark, A., et al. Training compute-optimal large language models. arXiv preprint arXiv:2203.15556, 2022.
Neumann, O. and Gros, C. Scaling laws for multi- agent reinforcement learning model. arXiv preprint arXiv:2210.00844 2022.
Padalkar, A., Pooley, A., Jain, A., Bewley, A., Herzog, A., Irpan, A., Khazatsky, A., Rai, A., Singh, A., Brohan, A., et al. Open x-embodiment: Robotic learning datasets and rt-x models. arXiv preprint arXiv:23100 20.223300.88844 2023.
Hu, A., Russell, L., Yeo, H., Murez, Z. Fedoseev, G., Kendall, A., Shotton, J., and Corrado, G. Gaia-l: A generative world model for autonomous driving. arXiv preprint arXiv:239.117080 2023.
Pearce, T. and Song. J. Reconciling kaplan and chinchilla scaling laws. TMLR, 2024.
Hussenot, L., Andrychowicz, M., Vincent, D., Dadashi, R., Raichuk, A., Ramos, S., Momchev, N., Girgin, S., Marinier, R., Stafiniak, L., al. Hyperparameter selec- tion for imitation learning. In International Conference on Machine Learning, pp. 4511-4522. PMLR, 2021.
Porian, T., Wortsman, M., Jitsev, J., Schmidt, L., and Carmon, Y. Resolving discrepancies in compute- optimal scaling of language models. arXiv preprint arXiv:2406. 19146, 2024.
Hutter, M. Learning curve theory. arXiv preprint MrXiv:2102.04074, 2021.
Raad, M. A., Ahuja, A., Barros, C., Besse, F,, Bolt, A., Bolton, A., Brownfield, B., Buttimore, G., Cant, M., Chakera, S., et al. Scaling instructable agents across many simulated worlds. arXiv preprint arXhv:2404.10179, 2024.
Isik, B., Ponomareva, N., Hazimeh, H., Paparas, D., Vas- silvitskii, S., and Koyejo, S. Scaling laws for down- stream task performance of large language models. arXiv preprint arXiv:2402.04177, 2024.
Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D. Scaling laws for neural language models. arXiv preprint arXiv:2001.0836) 2020.
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I., et al. Language models unsupervised multitask learners. OpenAl blog, 2019.
Reed, S., Zolna, K., Parisotto, E., Colmenarejo, S. G., Novikov, A., Barth-Maron, G., Gimenez, M., Sulsky, Y., Kay, J., Springenberg, J. T., et al. A generalist agent. arXiv preprint arXiv:2205.0611 2022.
Karpathy, A. NanoGPT. https: /ttp//ggttttttteee com/ karpathy /nanonan....... 2022.
Li, X., Hsu, K., Gu, Pertsch, K., Mees, O., Walke, H. R., Fu, C., Lunawat, I., Sieh, I., Kirmani, S., et al. Evaluat- ing real-world robot manipulation policies in simulation. arXiv preprint 24Xiv:2405.05941 2024.
Rosenfeld, J. S., Rosenfeld, A.. Belinkov, Y., and Shavit, N. A constructive prediction of the generalization error across scales. arXiv preprint arXiv: 0e... 2019.
Liu, Z., Mao, H., Wu, C.-Y., Feichtenhofer, C., Dar- rell, T., and Xie, S. A ConvNet for the 2020s, March 2022. URL http:/ 2201. 03545. arXiv:2201.03545 [es].
Sardana, N. and Frankle, J. Beyond chinchilla-optime
Accounting for inference in language model scaling laws. arXiv preprint arXiv:2401.0044 2023.
Maloney, A., Roberts, D. A., and Sully, J. A solvable model of neural scaling laws. arXiv preprint arXiv:2210.16859, 2022.
Sartor, S. and Thompson, N. Neural scaling laws for em- bodied ai. arXiv preprint arXiv:2405.14005 2024.
Silver, D., Schrittwieser, J., Simonyan, K., Antonoglou, I., Huang, A., Guez, A., Hubert, T., Baker, L., Lai, M., Bolton, A., et al. Mastering the game of go without human knowledge. nature, 550(7676):354-359 2017.
Micheli, V., Alonso, E., and Fleuret, Transform- ers are sample-efficient world models. arXiv preprint arXiv:2209.00) 00588, 2022.

---

Scaling Laws for Pre-training Agents and World Models
Su, H., Tian, Z., Shen, X., and Cai, X. Unraveling the mystery of scaling laws: Part i. arXiv preprint
arXiv:2403.06563, 2024.
Sutton, R. Reinforcement learning: An introduction. A Bradford Book, 2018.
Tian, K., Jiang, Y., Yuan, Z., Peng, B., and Wang, L. Visual autoregressive modeling: Scalable image generation via next-scale prediction. arXiv preprint arXiv:2404.02900 022002 2024.
Tuyls, J., Madeka, D., Torkkola, K., Foster, D., Narasimhan, K. R., and Kakade, S. M. Scaling laws for imitation learning in single-agent games. urXiv preprint, 2023.
Yang. M., Du, Y., Ghasemipour, K., Tompson, J., Schuur- mans, D., and Abbeel, P. Learning interactive real- real-world simulators. arXiv preprint 2r.iv:231000/// 14, 2023.
Yu, J., Li. X., Koh, J. Y., Zhang, H., Pang, R., Qin, J., Ku, A., Xu, Y., Baldridge, J., and Wu, Y. Vector-quantized image modeling with improved vggan. ICLR, 2022.
Zhu, Y. Kiros, R., Zemel, R., Salakhutdinov, R., Urta- sun, R., Torralba, A., and Fidler, S. Aligning books and movies: Towards story-like visual explanations by watch- ing movies and reading books. In The IEEE International Conference on Computer Vision (ICCV), December 2015.
12

---

Scaling Laws for Pre-training Agents and World Models
The appendix is organized follows.
Appendix contains details on the training of the model configurations, hyperparameters, and a description of the datasets used.
Appendix B contains results from Section 5.3.
Appendix C contains further details training world models robotics.
Appendix D contains further results demonstrating the link between pre-training loss and performance.
A. Scaling Experiments Further Details
This section provides experimental details for experiments on the primary Bleeding Edge dataset.
A.1. Hyperparameters
We trained two VQGANs from scratch with reconstruction losses.
BE-Small. Based on (Esser et al., 2021), uses d₂ = 256, V, = 4096, h = = 128, with 28M parameters, and a CNN design. It was trained on a single SkyGarden Bleeding Edge map.
BE--aaree Based on (Yu et al., 2022), uses d_ 540, V, 4096, 180, w = 300, with 150M parameters, and vision transformer design. was trained all seven Bleeding Edge maps.
We selected the numbers tokens per image based on qualitative assessment reconstructions. We found that 256 tokens per image was the minimum that still allowed reconstruction to capture the majority of salient gameplay details. However certain details still were lacking, such enemy player's health bars hence also considered 540 token version that provided higher quality reconstruction.
BC-CNN details. We use h = = 128. The 0.6M paramter CNN is similar to that used by (Baker et al., 2022), however it uses ConvNext blocks (Liu et al., 2022). The CNN produces an embedding of size 1024 which is then put through a linear layer to obtain vector matching the transformer's embedding dimension.
Transformer configurations are given in Table We describe the parameters for the WM-Token architecture. Note that MLP layers four times the width of embed dim. Model configurations roughly followed the model configurations used in Table A9 of (Hoffmann et al., 2022), where residual stream dimension, number of layers, and number of heads were roughly increased proportionally.
Table Transformer configurations. Here listed for tokenized architectures. Parameter count varies slightly BC-CNN inclusion of the embedding CNN and differing numbers of embedding parameters sizes.
N
Layers
Num heads
Embed dim
2M
3
3
180
4M
4
4
240
11
6
360
15M
4
512
27M
8
512
52M
10
10
640
110M
15
12
768
206M
16
16
1024
894M
23
14
1792
13

---

Scaling Laws for Pre-training Agents and World Models
A.2. Training Details
All transformers are trained with variant of nanoGPT (Karpathy, 2022) using PyTorch Lightning (Falcon & The Py PyTorch Lightning team, 2019).
This section lists key hyperparameters. Note that was important to find optimization settings that produced the lowest possible loss for a given model size. In general larger models require smaller learning rates. Our approach first optimized the smallest model through a grid sweep, we would then sequentially run a sweep over the next largest model, starting at the smaller model's optimized learning rate. Table 3-6 provide final settings.
Table 3. Hyperparameters for WM-Token with d2 =25 tokens per image observation. N
Seq len
Context length
Tokens per update
Learning rate
15M
10
2,720
522,240
0.0007
27M
10
2,720
522,240
0.0007
52M
10
2,720
522,240
0.0007
110M
10
2,720
522,240
0.0007
206M
10
2,720
522,240
0.00057
894M
10
2,720
2M
0.00028
Table 4. Hyperparameters for WM-Token with d, =540 tokens per image observation.
N
Seq len
Context length
Tokens per update
Learning rate
4M
10
5,560
533,760
0.005
11M
10
5,560
533,760
0.001
27M
10
5,560
533,760
0.001
52M
10
5,560
533,760
0.001
110M
10
5,560
533,760
0.0005
206M
10
5,560
533,760
0.0005
Table Hyperparameters for BC-Token with d, 550 tokens per image observation.
N
Seq len Context length
Tokens per update Learning rate
2M
10
5,560
533,760
0.0005
4M
10
5,560
533,760
0.0005
11M
10
5,560
533,760
0.0001
27M
10
5,560
533,760
0.0001
Table Hyperparameters for BC-CNN.
N
Seq len
Context length
Items per update
Learning rate
2M
10
10
2560
0.0003
3M
10
10
2560
0.0003
10M
10
2560
0.0003
26M
10
2560
0.0003
5IM
10
10
2560
0.0003
14

---

Scaling Laws for Pre-training Agents and World Models
Jeneey Sink
Hydrocorc
Aquadaats
Sky Garden
Lundside
Beneyard
Facility 665
0
12
3
f-5
1-
P8
Timestops
Figure 11. Example trajectories from dataset of 8.6 years of human gameplay in the video game Bleeding Edge across maps.
A.3. Dataset Details
Image observations were stored in MP4 format at 60fps, alongside binary files containing the associated controller actions. A time code extracted from the game was stored for each frame, to ensure actions and frames remained in sync at training time.
The Maps dataset comprised 60,986 matches, yielding 530,713 individual player trajectories (each around 9 minutes), totaling 27.8 7,88 TiB on disk. This amounted to around 8.6 years of gameplay. After downsampling to 10Hz (the frequency models trained on), this equated to 1.63B frames. This then divided into training validation / test sets by dividing the matches with 80: 10:10:10 split.
Our filtered Sky Garden dataset used the same 80:10:10 split and 10Hz downsampling, but focused on just one map, yielding 71,940 individual player trajectories, or 355.5M frames (around 1.12 years of game play).
For discretizing the controller actions, while the buttons are natively discrete, we discretize the and y values of the left and right joysticks into eleven buckets.
1. INFINITE DATA REGIME ALLOWED FLOPS
We wish study scaling in the infinite data regime, where training loss is not significantly effected by models repeatedly training on the same datapoints which can lead to overfitting effects. This section calculates the number of training tokens allowed for each model family trained in this work. Viewing Figure alongside these numbers confirms that models remain in the infinite data regime for all our experiments.
WM- -5ke BC- -5-Token-5500. We trained on the maps dataset, with 1.63B observation-action pairs. Models used the tokenized architecture with the large VQGAN, each observation-action pair creates 540 16 556 transformer inputs, 906B training tokens. (Muennighoff et al., 2024) observe that tokens may be reused up to four times with negligible departure from the infinite data regime. This produces 3.6T tokens. For 200M parameter model the compute allowed by the infinite data regime is C' 6ND = 6 x 200M x 3.6T 4.3 x 10 FLOPs.
for total of 1.63Bx556
WM-Token-225. This is trained on the Sky Garden dataset, with 355M observation-action pairs. Each pair is split into 256 + 16 = 272 tokens, for 97B training tokens, or 97B x 4 = 386B effective tokens. For a 200M parameter model the compute allowed by the 'infinite data regime' is C 6N D x 200M 386B 4.6 x 1020 FLOPPs
15

---

Scaling Laws for Pre-training Agents and World Models
BC-CNN. Trained maps dataset, but with token per observation-action pair, this creates possible 1.63B 4 6.52B effective tokens. A 50M parameter model uses C 6NT 6 x 50M x 6.52B 2.0 10 FLOPS.
B. Further Analysis Details
Experimental results supporting Section 5.3.
xx 10
Redt FLUh p0a
brtk 24
43M poreo
Optimul parameeeess
Optimal tokens
10
Loss
2x 10
10²
10
10
10"
10"
10"
10"
10"
10"
10"
10"
10"
FLOPs
FLOPs
FLOPs
Figure Relating Section 5.3, character-level (low compression). Utilising the frontier (middle and right) we derive power cnefficient for Noytimal 0.66 and Dopimat as 0.34.
AXIM
mocdried 11,P Puece lay,
Bormad FLUTS
L6G pains
4x 10
Towe b-g
6x 10
parumelers
3x 10
Optimal tokens
10
Loss
10
Optimal
10
4x 10
6x 10
10
ELOPs
FLOPs
FLOPs
Figure Relating Section 5.3, GPT-2 tokenizer (high compression). Utilising the frontier (middle and right) we derive the power law coefficient for Nogimal as 0.44 and D Depimal as 0.56, an increase from 0.66 in Figure 12 found when utilising a lower compression character-level tokenizer.
16

---

Scaling Laws for Pre-training Agents and World Models
C. World Modeling for Robotics Experimental Details
This section provides experimental details for WM experiments on the secondary RT-1 dataset.
C.1. Dataset
We resized the RT-1 dataset to 1 128x128 pixels per image. For action labels, we take the 3D corld_vectoo coordinates, combined with the ID aaripper_closedneeee closedness. actionn vector, to make an action vector with four dimensions. All are in the range- -1 to and these are discretized into 500 evenly spaced buckets.
C.2. VQVAES
We trained a set of five VQVAEs using the implementation from httpp: com/ httm//adavbh12///0 VAE. We set 2g [16, 36, 64, 100, 256] and Vo 4096, training each VQVAE for 40,000 updates on batches of 128. Reconstructions visualized in Figure 14.
Reconstructions
Ground truth
256 tokens
100 tokens
64 tokens
36 tokens
16 tokens
Figure 14. VQVAE reconstructions on the RT-1 dataset for differing numbers of tokens per observation, E [16, 36, 64, 100, 256].
C.3. Transformer training details
Table 7 provides training details for the model sizes tested. Figure shows one example set of training curves per VQVAE.
Table Hyperparameters for WM-Token RT-1 experiments.
N
Seq len
Context length
Tokens per update
Learning rate
0.08M
2
2(2. + 4)
34,000
0.01
0.2M
2
2(20 + 4)
34,000
0.005
0.28M
2
2(20 + 4)
34,000
0.004
0.54M
2
2(20 +
34,000
0.0027
0.99M
2
2(20 +
34,000
0.002
17

---

Scaling Laws for Pre-training Agents and World Models
16, Noptinal & C0.56 Noptinal & DO.44
20
10
thea ewest FLAPs Poeur lan wrd
For roald FLOPS Pouer and brg 44
822 parco
4x 10"
Optimall puramelen
Optimal tokens
Loss
10
1x
10
2 10"
10"
10"
10"
FLOPs
FLOPs
FLOPs
20
36, Noptimal
60 N Noptimal D0.40
puaca
teei pres 1LAh
10
10
Opt Optimal paramelan
Optimal tokens
Loss
Tuseatect
10
3x 10"
10²
2x
10
10
10
10
10
FLOPs
FLOPs
FLOPs
64,
Noptimal x 0
Noptimal x D0.39
Ibe meer 11,0 Pueae a (61)
Bormal FLN betk
4x 10
10a
Optimall parameless
Optimal tokens
Loss
30x10
10
m
10
2 10
10t5
10
10
10
10
10
FLOPSs
FLOPs
FLOPs
100, Noptimal x 00 N Noptimal x D0.40
20
modd grea
Bei roall yoe FlOPs
10a
4x 10
Optimnal par paramelees
Optimal tokens
1i frales
Loss
10
3x
10
10
10
10
10
10
10"
10
FL.0Ps
FLOP
256,
N Noptmal x 0.66
N Noptiaal x 0..34
3x 10
6234
10
amaid
Plok
3x 10
Noer
1043
Optimal paramelees
10
Optimal tokens
Loss
mare
3x
10'
10
10
10
10
10
10"
10"
FLOPSs
FLOPS
FLOOPs
Figure RT-1 experiments. Note the optimal parameter coefficient increases number of tokens per observation.
18

---

Scaling Laws for Pre-training Agents and World Models
Efficient fronteer method
0.64
a where N x Ca
0.62
0.60
0.58
0.56
10 2
Tokens per image
Figure 16. RT-1 experiments. Optimal parameter coefficient vs. number of tokens per observation, with three repeated runs per VQVAE, Correlation, A = 0.61.
19

---

Scaling Laws for Pre-training Agents and World Models
D. Pre-training Loss vs. World Modeling Metrics
This section presents evidence for pre-training loss correlating with WM performance. We use metrics commonly used to assess the quality of the world models (Yang et al., 2023), originally developed the video generation literature. Conditioned on an initial real frame and a sequence real actions, we compare the observations generated by world model, with the real sequence of observations, measuring FVD and LPIPS. Specifically, generate 1024 videos each of 10 seconds. We perform this for various checkpoints on each size in our WM-Token-256 set of models. This allows plot of the checkpoint pre-training loss vs video generation metric to be assessed.
Figure shows results. We find correlations of 0.77, 0.83 for LPIPS and FVD respectively. Two early checkpoints from the 894M model are only significant anomalies trend of metrics improving with loss. This evidences the strong relationship between pre-training loss and world model quality.
10'
Size: 27M
Size: 27M
Sizc: S2M
Size: S2M
Size:
Size: 206M
LPPDS
Size: 205M
10
Size: 594M
Size: 894M
10
Correlation, R 0.83
Correlation, R 0.77
2.25
2.50
2.55
Loss
Figure Our experiments suggest pre-training loss is good proxy generation quality of world model, FVD LPIPS.
20

---

Scaling Laws for Pre-training Agents and World Models
E. Additional Small Scale BC-Token Experiment
This section presents an additional set of results for a set of BC-Token models smaller parameter counts. The aim had been enable coefficient estimation using the frontier method. However, even smaller sizes we did not observe the overlap required for this method presented in Figure 18. Using the parametric fit method we observed, Nooptimal C Napumal D0.65 which are in line with the coefficients of larger BC-Token-5544 models in Table 1.
7 x 10-1
147k params
310k params
647k params
1M params
Parametric fit
6 x 10-1
Loss
5 10
10
10 FLOPs
10
Figure Additional of results for a set Boken models of smaller parameter counts.
Using parametric method observed.
Noqtima O 0.33 Noptiaa D 0.65
Table 8. Transformer configurations for small scale BC-Token experiments.
N
Layers
Num heads Embed dim
147k
25
310k
46
647k
86
1M
2
120
Table 9. Hyperparameters for BC-Token with d, =540 tokens per image observation. scale experiment.
N
Seq len
Context length Tokens per update Learning rate
147k
5,560
500,000
0.0005
310k
5,560
500,000
0.0005
647k
10
5,560
500,000
0.0005
IM
10
5,560
500,000
0.0005
21


Reviews:
[{"review_id": "bfXC174jfV", "reviewer": "Reviewer_TxDt", "summary": "The paper explores scaling laws in embodied AI, especially for the pre-training stage of world models and agent behavior. The authors show that power laws similar to those in LLMs are observed in world modeling and behavior cloning, but with coefficients influenced by the tokenizer, task, and architecture. The study provides insights into optimal model and dataset sizing for these tasks.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review, we are pleased your judgement of the paper comes out on the side of acceptance and agree that your suggestions would further improve the paper. Due to logistical constraints, we are not able to complete all of these requests, but commit to completing the smaller scales BC-Token experiments (which we agree is an excellent idea), before any camera-ready deadline. Below we respond to your three primary questions, and several further key comments we noted.\n\n- __Q1.__ Additional evidence … to show the correlation between the training loss and game performance\n\nWe have done our best to frame all claims in the paper as optimizing for loss in the pre-training phase of agents, rather than for downstream performance, as these are the concrete results we have. We have further provided conceptual arguments and experimental evidence on the link between pre-training loss and performance but agree this should continue to be explored through future work. \nDoing this effectively for our dataset, environment and setup, we feel we would require resources beyond the scope of this paper. We would require implementing a form of post-training phase to maximize some reward signal, and engineer an automated, distributed solution to enable rollouts in this game at large scale (something Bleeding Edge does not natively do).\nAs such, similar to the LLM reference noted by the reviewer (their [2]), we believe this is best left to a separate paper. \n\n- __Q2.__ Smaller models in the BC-Token experiments\n\nThank you for this very sensible suggestion! Our original protocol had aimed for consistent model sizes across tasks and tokenizers, but on reflection, we agree that smaller models could be used for the BC-Token experiments. Logistically this will be difficult to complete within the rebuttal period (very small models do not fully utilize GPU FLOPs and hence are very time inefficient), but we are happy to commit to completing this before any camera-ready deadline (should we be fortunate enough to have the paper accepted).\n\n- __Q3.__ Additional extrapolation experiments\n\nWe agree with the spirit of this request – the gold standard test for scaling laws is how well they predict orders of magnitude out. As the reviewer mentions, this would require more compute than used in our current paper, which is a hard for us to commit to. We hope our current results, which span around three orders of magnitude, are sufficiently interesting for much of the embodied AI community working which has not yet advanced to the model sizes in LLMs.\n\n- __Comment 1.__ Claims feel stronger than results … since the scaling laws aren’t really general and are influenced by the specific task, tokenizer and architecture\n\nPlease let us know if there are specific wording changes you’d recommend. We aimed to write the paper in a way that would engage researchers across embodied AI, whilst avoiding overclaim or hype. Power laws did consistently emerge across the two domains tested, and one of our main contributions is that task, tokenizer and architecture _do_ impact coefficients. This is an important departure from LLM research, where datasets, tokenizers and architectures usually only see minor variations.\n\n- __Comment 2.__ Dataset is proprietary, reproducing these results are substantial, limiting accessibility for researchers with fewer resources\n\nThe Bleeding Edge dataset was accessed under a data sharing agreement and unfortunately, we do not have agency to share ourselves. On the other hand, the RT-1 experiments are fully reproducible. This follows recent high-impact embodied AI work (Genie: Generative Interactive Environments) which described a lightweight environment for other researchers to experiment in. Note our LLM experiments are also designed to help researchers capture the essence of challenges we identified in a more accessible modality.", "claims_and_evidence": "* Scaling laws similar to those in LLMs can be observed in world modeling with tokenized observations and actions → Partially insufficient \n  - The authors show extensive experiments to prove the scaling laws for training loss, but only minimal examples of how this training loss translates to real world performance on embodied AI tasks. Additional tasks to show the training loss as a proxy would make this claim stronger.\n* The optimal trade-off between model and dataset size in world modeling is influenced by the tokenizer’s compression rate → Sufficient\n  - Two examples of tokenizers are used to support this claim. Additional examples using the small RT-1 dataset also support this claim, but the sizes used are quite small due to the limited size of the dataset.\n* Scaling laws for BC with tokenized observations are harder to observe under modest compute budgets. The optimal trade-off favors smaller models and more data → Partially insufficient\n  - The authors claim that models with size > 2M params don’t saturate over the flops range considered and use the parametric fit instead of the frontier fit method. Can additional models with smaller sizes <2M be used to show that the parametric fit and frontier fit curves match for smaller sizes?\n* Scaling laws similar to those in LLMs can once again be observed in BC with one continuous encoding per observation → Sufficient\n  - The power law curves for training loss with BC-CNN show a strong trend as claimed.", "methods_and_evaluation_criteria": "The benchmark dataset used from the game Bleeding Edge is a reasonable choice for the proposed claims, but might not be sufficient to show that the claims also hold for other examples of embodied AI - especially real world tasks with a long tail of difficult scenarios. The authors show that the training loss has a strong correlation with the final task performance for the few tasks studied, but it’s not clear if these will also hold for other general cases of BC and world modeling (like discussed in [1]), as is even discussed for LLMs in [2]. \n\nAlso, previous studies like [3,4] have shown significant impact of dataset quality on downstream performance for agents. Also, it would be very useful if the 7 Maps dataset is available publicly to allow reproducibility of this research.\n\n[1] Codevilla, Felipe, Eder Santana, Antonio M. López, and Adrien Gaidon. \"Exploring the limitations of behavior cloning for autonomous driving.\"\n\n[2] Liu, Hong, Sang Michael Xie, Zhiyuan Li, and Tengyu Ma. \"Same pre-training loss, better downstream: Implicit bias matters for language models.\"\n\n[3] Bronstein, Eli, Sirish Srinivasan, Supratik Paul, Aman Sinha, Matthew O’Kelly, Payam Nikdel, and Shimon Whiteson. \"Embedding synthetic off-policy experience for autonomous driving via zero-shot curricula.\" \n\n[4] Belkhale, Suneel, Yuchen Cui, and Dorsa Sadigh. \"Data quality in imitation learning.\"", "theoretical_claims": "N/A all the claims in the paper are empirical and supported by experimental results.", "experimental_designs_or_analyses": "The experiments are well designed and detailed to support the claims for the proposed domain of agents and world modelling for the Bleeding Edge game. Additional evidence of final game performance of the agent and world model using the proposed pre-training setup to show the strong correlation between the training loss and the game performance will help solidify the claims further. \n\nThe extension of these claims however to the entire embodied AI domain for agents and world models aren’t really conclusive - especially because the variation between optimal dataset and model sizes varies quite a lot between the different tasks and tokenizer.", "relation_to_broader_scientific_literature": "The contributions of the paper, especially showing that the scaling laws from LLMs extend to embodied AI are significant, and will accelerate the development of models for the community by providing suggestions for optimal model and dataset sizes. The claims however feel stronger than the actual results, especially since the scaling laws aren’t really general and influenced by the specific task, tokenizer and architecture.", "essential_references_not_discussed": "None that are obvious or very well known.", "other_strengths_and_weaknesses": "Strengths\n- The authors successfully demonstrate that power laws similar to those in language models also apply to world modeling and behavior cloning. This extension is crucial as it guides researchers in optimizing resource allocation for embodied AI tasks.\n- The paper employs a rigorous methodology, using clear definitions and complementary fitting methods to establish scaling laws.\n- Key findings include the influence of tokenizer compression rate on optimal model size in world modeling and the preference for smaller models with more data in behavior cloning with tokenized observations.\n\nLimitations:\n- While the paper justifies using pre-training loss as a proxy for performance, this approach has limitations. Studies have shown mixed results regarding the correlation between pre-training loss and downstream performance.\n- The paper does not fully address domain-specific challenges such as physical interaction complexity and the lack of long tail scenarios.\n- The curves generated using parametric fit definitely support a stronger claim than those with the parametric fit.\n- The paper could benefit from more detailed analysis of dataset diversity and potential biases.\n- The dataset used for the experiments is proprietary, as well as the computational requirements for reproducing these results are substantial, limiting accessibility for researchers with fewer resources.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "1. Additional evidence to support the claim that the pre-training loss correlates strongly with the final performance metrics will definitely be useful. Especially because the domain of embodied AI with agent modelling suffers from error accumulation over steps, resulting in out-of-distribution states and sub-optimal performance in the long tail of scenarios.  \n2. The authors claim that the scaling laws for BC with tokenized observations are hard to observe under the current compute budgets. Could additional evidence using smaller models or a different dataset be added to show that the scaling laws actually satisfy the authors’ claims?\n3. Figure 9 shows the extrapolation of the scaling laws to a much larger model - additional examples of large scale models for the other claims will also help strengthen that they are satisfied over a wide range of model scales. (I do understand that this could require significant training time, so will only be a positive if it can be added)", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_concerns": "N/A"}}, {"review_id": "KJSyFjZfdb", "reviewer": "Reviewer_Pwoy", "summary": "This paper investigates the existence and characteristics of scaling laws in embodied AI tasks, specifically world modeling (WM) and behavior cloning (BC), drawing parallels to scaling laws observed in large language models (LLMs). Through extensive experiments on large-scale video game datasets (e.g., Bleeding Edge) and robotics data (RT-1), the authors demonstrate that power-law relationships between model size, dataset size, compute, and pre-training loss also apply to WM and BC. Key findings include:\n\n- Scaling laws for WM are influenced by tokenizer compression rates. Higher compression rates lead to more reliance on dataset size.\n\n- BC with tokenized observations requires prioritizing dataset scaling under modest compute budgets, while BC with CNN-based architectures favors model scaling.\n\n- Small-scale language modeling experiments validate mechanisms behind observed scaling phenomena (e.g., sparse supervision and target granularity).", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. We are delighted to have successfully communicated the value of our study. We agree that extending principles from LLMs to embodied AI tasks is a timely and important avenue of research. Allow us to respond to your comments below.\n\n- __Comment 1.__ RT-1 dataset has limited diversity and is small scale. How do conclusions apply to tasks where generalization is important, and larger scale datasets.\n\nWe recognize that the RT-1 tasks, policies and visuals contain only limited diversity. On the other hand, the Bleeding Edge dataset contains a huge amount of diversity, with a pool of seven maps, tens of thousands of demonstrators of differing skill levels.\nWhile these two datasets are from differing domains, we did not observe contradictions between the two sets of experiments. We hope that analyzing the results in combination will give researchers enough evidence that similar patterns would also emerge in a large diverse robotics dataset. \n\n- __Comment 2.__ Effects of tokenizer compression rates and target granularity lack theoretical grounding\n\nThank you for this comment. This is one of the directions for future work we are also most excited about, and have started very initial experiments in this direction. We sketch our current thinking on this issue in case of interest.\n\nWe have begun considering the case of lossless compression, where more-tokens-per observation (lower compression) should make next-token-prediction easier, both in the sense of the pattern being simpler (less parameters required) and less stochastic noise (less data required). We suspect this leads to lower values of both alpha and beta. But the key question is which shrinks faster (since their ratio is the thing that matters). Which might be harder to reason about without further assumptions.\n\nOne thing that makes analysis tricky, is that losses may not be consistent across tokenizers even for the lossless case, and in reality, tokenizers become more lossy at higher compression rates.\n\nRegarding this paper submission, we believe that a deeper analysis of compression rates is best left to a separate paper, and perhaps domains more straightforward than embodied AI. Should you be interested in exploring this direction further, please reach out following the review period.", "claims_and_evidence": "All claims are supported by thorough experiments that are carefully designed.", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "All experimental are carefully designed and conducted.", "relation_to_broader_scientific_literature": "Scaling laws play a crucial role in the development of large AI models, such as LLMs. This paper makes significant contributions to the understanding of scaling laws in agent learning by utilizing expressive architectures. These insights will be instrumental in unifying scaling laws across various domains", "essential_references_not_discussed": "To the best of my knowledge, all related works are cited/discussed.", "other_strengths_and_weaknesses": "Strengths:\n\n- The work bridges a critical gap in understanding scaling laws for embodied AI, extending principles from LLMs to WM and BC. This is timely, given the growing interest in scaling generative models for robotics and interactive agents.\n\n- The paper validates findings across diverse datasets (video games, robotics) and architectures (tokenized vs. CNN-based). The inclusion of small-scale language experiments to explain mechanisms adds methodological depth.\n\n- The analysis of tokenizer compression rates and architecture choices provides actionable guidelines for optimizing compute allocation in real-world applications.\n\n- I really like the meta-analysis linking pre-training loss to downstream metrics (e.g., FVD/LPIPS for WM), which strengthens the case for using pre-training loss as a proxy in scaling studies.\n\nWeaknesses:\n\n- The existing datasets for robotics lack sufficient diversity, which may limit the value of the conclusions presented in the paper, particularly when downstream tasks emphasize generalization capabilities [1]. The RT-1 robotics dataset’s small scale also raises questions about extrapolation to large-scale datasets in robotics like OXE dataset.\n\n- The observed effects of tokenizer compression rates and target granularity lack theoretical grounding. While empirical results are compelling, deeper mechanistic explanations would strengthen the work.\n\n\n\n[1] A Taxonomy for Evaluating Generalist Robot Policies", "other_comments_or_suggestions": "N/A", "questions_for_authors": "N/A", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "S6Q0La40Xv", "reviewer": "Reviewer_pTPD", "summary": "This paper investigates the scaling laws in embodied AI. Specifically, this paper focuses on the infinite data regime and generative pre-training objectives, which include behavior cloning and world modeling. The scaling laws are observed in the following two cases through the experiments. One is world modeling with tokenized observations and actions while the other is behavior cloning with one continuous encoding per observation. Experiments and analysis are conducted in the domains of video games and robotics.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to review our paper. Please see below our responses.\n\n- __Comment 1.__ Infinite data regime assumption\n\nThank you for drawing attention to this important detail. To clarify, all experiments in our paper for both domains (video games and robotics) are conducted in the infinite data regime (Section A.3.1 computes the amount of FLOPs that would violate this assumption in our set up for the main experiments). \n\nInvestigation into scaling laws outside the infinite regime is an active research area in LLM research (Scaling Data-Constrained Language Models). As we understand, the current thinking is that scaling laws do still exist in this setting, but are influenced in adverse ways by repeated epochs.\n\nWe believe the reviewer is correct to call out this as a line of research very relevant to embodied AI, which typically falls within the data-constrained regime. We believe it falls outside the scope of a first investigation into scaling laws in embodied AI, for which we have chosen to focus our investigation on the effect of tokenizer, task and architecture. \n\n- __Comment 2.__ Pre-training loss and downstream performance\n\nWe have done our best to frame all claims in the paper as optimizing for loss in the pre-training phase of agents, rather than for downstream performance, as these are the concrete results we have. We have further provided conceptual arguments and experimental evidence on the link between pre-training loss and performance. Note the Nethack and Atari experiments, while being of limited complexity, are not necessarily small-scale (up to 10^18 FLOPs).\n\nProviding further evidence effectively for our dataset, environment and setup, we feel would require resources beyond the scope of this paper. We would require implementing a form of post-training phase to maximize some reward signal, and engineer an automated, distributed solution to enable rollouts in this game at large scale (something Bleeding Edge does not natively do).\n\nIn general, we agree with the reviewer that the link between pre-training loss and online performance is important enough to deserve further study, but leave this to a separate investigation.\n\n- __Other comments__\n\nThank you for these suggestions. Should we be fortunate enough to have the paper accepted, we will take your suggestion and use the allowed extra page to enlarge figures for clarity. Regarding the paper organization, please confirm whether other reviewers support your proposed changes. We would then be happy to reorder the paper as you suggest in the next version.", "claims_and_evidence": "My primary concern is that this paper is built upon the assumption of an infinite data regime, where samples are not trained more than once. While this assumption holds for LLM on NLP pre-training, is it applicable in the cases used in this paper, such as specific video games or robots in a fixed environment?", "methods_and_evaluation_criteria": "This paper claims the pre-training loss could be used to represent the downstream performance and thus serve as the indicator for scaling laws. However, this connection is simply obtained from small-scale environments like Atari and NetHack. Extending this observation to more complex scenarios (e.g., video games and robotics) has no clear evidence or guarantee. Hence, whether the observed scaling laws can represent the real robot performance remains unclear.", "theoretical_claims": "No theoretical claims are provided in the main paper.", "experimental_designs_or_analyses": "This paper only uses generative pre-training loss for evaluation. How this proxy reflects real performance is lacking.", "relation_to_broader_scientific_literature": "Compared with previous works, this paper chooses to analyze the scaling laws of embodied AI by generative pre-training loss in the infinite data regime.", "essential_references_not_discussed": "To the best of my knowledge, the references are sufficiently covered.", "other_strengths_and_weaknesses": "1. The overall paper is not well organized. For example, all figures in this paper are too small.\n2. The subsections in this paper are also not well structured. For example, it is better to put Sec. 2.2 into section 3 due to Sec. 2.2 can be viewed as the foundation of the method. Also, putting Sec. 3.3 to experiment section would let the reader easy to follow. Overall, the paper flow and the figure layout are suggested to be improved.", "other_comments_or_suggestions": "Please refer to the issues raised above.", "questions_for_authors": "Please refer to the issues raised above.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ER5t7QrbYY", "reviewer": "Reviewer_8UYY", "summary": "The work presents a scaling law study examining the behavior of action and observation prediction models (behavior cloning and world models, respectively). The main results characterize trade-offs between model and dataset scaling given a fixed compute (FLOPS) budget. One result evaluates world model prediction performance based on the size of the discrete token vocabulary for observations. A second result evaluates behavior cloning performance based on using discrete or continuous observation embeddings. Additional analyses justify the choice of next-token loss for downstream task performance and use other domains (text and robotics) to examine potential causes for the different scaling behaviors.\n\nThe main findings are that:\n- Increasing the discrete token vocabulary of a world models leads to better scaling in model size compared to data.\n- Discrete observations in behavior cloning favor smaller models with more data compared to continuous observations. Scaling laws like in LLMs are less clear (under the experimental FLOP budgets) for discrete observations.\n\n\n## update after rebuttal\n\nThe replies addressed my questions, but my score was already quite positive. No changes.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Many thanks for taking the time to review our paper in such detail, and raising several points we had not considered. We are pleased to see the value of our contributions has been recognized. Below, we respond to your main questions, followed by your ‘other comments’ and finally several points raised elsewhere in the review. \n\n- __Q1.__ Discrete vs continuous encoders for BC.\n\nThanks for this comment – we appreciate the reviewer’s insight into trying to reconcile the two plots, which is something we agree is worth understanding in more depth. First, as discussed in __Other comment 1.__ caution should be applied as the two BC losses are in fact subtly different, e.g. one should not say the BC-CNN loss is ~10x higher than BC-Token’s. \nHowever, the fact BC-CNN asymptotes quicker, means those models are approaching the natural entropy in the data ($E$ in Eq. 7) at modest compute budgets, while BC-Token does not show signs of this. One hypothesis is that BC-CNN models are therefore more compute efficient, though an alternative hypothesis is that differences in loss modelling cause this.\nWe propose to investigate this further when training the new smaller BC-Token models requested by Reviewer TxDt, by computing the equivalent BC-CNN loss for the new models, to better understand how these two losses align. Again, thanks for the suggestion.\n\n- __Q2.__ Parametric fit distorting results\n\nWe assume the reviewer is referring to the BC-Token experiment? We expect that following Reviewer TxDt’s recommendation to train a set of smaller models to encourage saturation, we will be able to provide coefficients for the frontier fit method resolving this issue.\n\n- __Q3.__ Inference speed and model size\n\nWe can confirm all BC-Token and BC-CNN models considered in our experiments can be run at >10Hz on a V100 GPU. We do not currently have numbers to hand on how inference speed changes with model sizes.\n\n- __Q4.__ Figure 11 correlation\n\nThe computed correlation comes out as 0.61, which is considered a moderate strength relationship. We will add this to the figure caption. The high variance comes since a set of models are trained for each point – one model in the family failing to converge properly affects the coefficient entirely. This was something that occurred more often for models trained on larger tokens-per-image. We believe with further optimizing of training schedule these models could converge more consistently and the pattern would emerge even stronger. \n\n- __Q5.__ Investigating effect of tokenization of actions on coefficients\n\nThis is a nice suggestion we had not thought of! Does the action tokenizer effect coefficients in similar ways to image tokenizers? Whilst most action dims are natively discretized (buttons), the natural lever to play with here would be number of bins for the continuous joystick dims. But we could also consider grouping together various button combinations to require less tokens per action prediction. We feel this goes a little beyond the current scope of the paper but would be excited to see investigation of this in follow up work.\n\n- __Other comment 1.__ Figure 1 with differing loss scales\n\nWe apologize for not making this clearer. In fact, the loss magnitudes are not directly comparable between BC-Token-540 and BC-CNN. Each model optimizes a subtly different losses; BC-Token-540 predicts action dimension-by-dimension, while BC-CNN produces predictions for each action dimension independently (Section 3.2). Our implementations also aggregate the losses in different ways. We will edit the figure caption to make this clearer.\n\n- __Other comment 2.__ Small figures\n\nApologies, we will expand figure sizes in the next version.\n\n- __Other comment 3.__ Line 370 unclear\n\nApologies for this, you were correct in your guess. We will clarify.\n\n- __Comment 1.__ Re-analysis of BC ... unsure about infinite data assumption\n\nTo confirm, Tuyls et al. used pre-trained policies so could generate unlimited data and remain in the infinite data regime.\n\n- __Comment 2.__ Character prediction experiments used parametric fit\n\nWe agree frontier loss would be preferable. While we could have used frontier fit for the Dense loss, this did not appear possible for spare loss super-classed (Figure 10). We favored consistency in our fitting method across these experiments.\n\n- __Comment 3.__ Conclusions about scaling hard to understand\nThanks for this comment. We have focused on presenting nuanced details from the scaling analysis and have so far not offered clear conclusions for practitioners. Would the reviewer agree with the below concrete recommendations?\n1) When training BC-Token models, model sizes should be substantially smaller than for BC-CNN.\n2) For WM-Token models, increasing tokens-per-image of the tokenizer should be done in parallel with increasing model sizes.\n3) As per response to __Q1.__ we may be able to advise a preference for BC-CNN models, but await our analysis contrasting BC-CNN and BC-Token losses.", "claims_and_evidence": "The claims and their evidence:\n- claim: World models show power law scaling.\n\t- evidence: Figures 1, 5, 6 and the associated experiments training a world model on Bleeding Edge data.\n- claim: Imitation learning models show power law scaling.\n\t- evidence: Figures 1, 7, 8 and the associated experiments training behavior cloning on Bleeding Edge data.\n- claim: Tokenized imitation learning models favor smaller models with more data.\n\t- evidence: Table 1 on the frontier fit analyses (and parametric fit for BC-Token-540). Also Figure 7.\n- claim: Continuous observations for imitation learning models favor larger models.\n\t- evidence: Table 1 on the frontier fit analyses (and parametric fit for BC-Token-540). Also Figure 8.\n- claim: The trade-off between model and dataset size for world models is correlated with the number of tokens per observation.\n\t- evidence: For the base experiments Table 1 and Figures 5 and 6.\n\t- evidence: A test on RT-1 scaling training varied tokenization sizes and world models (Figure 11).\n- claim: Next token prediction loss is a good proxy for reward achieved by behavior cloning. \n\t- evidence: Figure 2, which re-analyzes a previous study on model scaling for behavior cloning. This shows correlation coefficients in the range of -0.94 to -0.98.\n- claim: Next token prediction loss is a good proxy for image reconstruction quality achieved by a world model.\n\t- evidence: Figures 3 and 7, which analyze image reconstruction quality (FVD and LPIPS) compared to world model size. These show correlation coefficients of 0.83 and 0.77.\n- claim: The amount of token-level supervision and label super-classing explain the shallower scaling behavior of behavior cloning models compared to world models.\n\t- evidence: Shakespeare text prediction scaling analyses that manipulate prediction tasks and class aggregation.", "methods_and_evaluation_criteria": "Yes. Scaling on game data in an \"infinite data regime\" is appropriate to measure scaling laws. The main body and supplement provide evidence that the loss in this regime is strongly correlated with reward for behavior cloning (correlation coefficient larger magnitude than 0.9) and for observation prediction (aka world modeling, correlation coefficient 0.83 for FVD and 0.77 for LPIPS). While the evidence is from different game environments, it provides reasonable grounds to extrapolate.\n\nIdeally the games used would be more diverse (than maps from a single environment), but that is not strictly needed for this type of analysis. It may alter the scaling coefficients observed due to the heterogeneity in observations and actions, but it would be surprising to learn that other games violated these patterns in the main (and worth separate study).\n\nThe supplemental analyses of Shakespeare text and RT-1 observations add support to the core scaling claims, though I have some modest concerns about the methods and how well they would translate across model tasks (see below).", "theoretical_claims": "No proofs are in the paper.", "experimental_designs_or_analyses": "Re-analysis of behavior cloning scaling in prior games. No obvious issues, though I am not familiar with the original paper and not sure if the assumption of an infinite data regime matches.\n\nThe scaling analyses for behavior cloning and world modeling in Bleeding Edge. The questions below addresses specific points. Most pressing is how to interpret the BC comparison given the methodological fitting differences between the discrete and continuous cases. It is not clear what conclusions we can draw from differences in the methods and given the lack of loss saturation given the FLOPS budget for the discrete case.\n\nWorld model fitting extrapolation. No obvious issues, but it returns to the concern about drawing conclusions in the case where compute requirements prevented full testing.\n\nCharacter prediction analyses of proposed mechanisms for different scaling behavior. These were a useful supplement to add evidence in favor of the hypotheses around lack of supervision and token granularity to explain the WM vs BC differences. My only (minor) concern is methodology: the fits are all estimated using parametric fit, despite preferring frontier fit for most of the other results in the paper. This is particularly true since the BC-CNN scaling results seem sensitive to this choice.\n\nRT-1 observation encoding results to explain world model scaling behavior in model parameters compared to tokenization vocabulary size. These are helpful to have for more context, but raise questions about how strong the relationship is.", "relation_to_broader_scientific_literature": "The prior literature on scaling laws and methodological improvements are connected to the choice of methods in this paper. The results in the paper are contrasted to prior work on auto-regressive modeling in video and images, where the newer methodology contradicts past findings. The paper also discusses prior scaling law research in the embodied context, including the work that is re-analyzed for behavior cloning.", "essential_references_not_discussed": "No references were missing that I consider essential.", "other_strengths_and_weaknesses": "# strengths\n- clarity: Explains the methodological choices made and their application that helps introduce readers less familiar with the technical aspects of the neural scaling law literature.\n- rigor: The main claims are supported by convergent lines of evidence aside the main results. This includes the correlation studies for BC and WM with downstream task performance and follow-up experiments around the scaling variation hypotheses in the Shakespeare and RT-1 domains.\n- significance: Establishing trade-offs in scaling architectures for embodied AI (at least in games) is important as these techniques see wider adoption. The methodological rigor and references here will help that as well.\n\n\n# weaknesses\n- clarity: The conclusions about WM scaling are hard to understand. The meaning is obstructed by jargon. There are many ideas being packed in, but it's hard to understand the overall narrative.\n\t- It would help to make some statement along the lines of \"When data is limited, favor scaling X.\" \"When model size is constrained, favor scaling Y.\" for each case of: the WM tokenizer vocabulary size, BC tokenizer method.\n\t- Similarly, the BC scaling conclusions are difficult to understand as stated.", "other_comments_or_suggestions": "I've included any non-critical questions in this section.\n\n- Figure 1 can be misleading around BC scaling. Multiple times I've misread the figure due to the magnitude differences of losses for the tokenized models (roughly 0.4 to 0.5) and continuous models (3.5 to 4.0). It may help to more explicitly call this out or somehow make it more apparent that the CNN losses are roughly an order of magnitude higher. This difference is important when considering the practical implications of the work for downstream tasks.\n\t- More generally, it may help to provide some remarks on what this implies for embodied AI training.\n\t- What are the implied best practices from these results?\n\t- It seems that the conclusion is training a discrete model is best with as much data as possible is best (modulo saturation concerns for a given compute budget).\n\t- Under what conditions of data limitation (for a given compute budget) does it become beneficial to use a continuous model?\n- Figures 2 through 10 are too tiny to read. It would help to make the key figures larger and move less crucial points to the appendix.\n\t- Perhaps the figures justifying the correlation of infinite data loss to outcomes.\n- \"This matches the magnitude of decrease seen in Table 1 from 0.66 to 0.32, indicating that the proposed mechanisms explain our findings.\" (line 370)\n\t- This was confusing as stated. After some looking I believe it's intending to reference C^a for BC-CNN under Frontier Fit and BC-Token-540 under Parametric Fit.", "questions_for_authors": "- [Q1] What is the conclusion to draw for image encoders being discrete vs continuous?\n\t- The CNN (continuous) model has higher loss and asymptotes being more shallow. Is this connected to the size of the latent representation?\n\t- The tokenized (discrete) model has not started plateauing in the same way.\n- [Q2] Is the use of parametric fit distorting these results at the highest FLOPS due to models not saturating? Did the discretization have too many tokens?\n- [Q3] How does inference speed scale with model size?\n\t- This is not crucial, but if the data was already on hand it would strengthen the results.\n\t- Specifically this seems important in the behavior cloning case, where the learned model would need to run inference at roughly 10 Hz (to keep up with a game).\n- [Q4] What is the correlation coefficient for Figure 11? \n\t- The association looks weak to the eye. And there is also a potentially small effect size given the very compressed scale.\n- [Q5] Is there any analysis on the discretized vocabulary size of BC that would be comparable to the world model tokenization size experiments?\n\t- Given the lack of saturation and the compression ratio results, it would be helpful to have a similar parallel analysis for the BC case. This is a \"nice to have\" that would strengthen the overall narrative.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper introduces the first scaling laws for embodied agents, concretely behavior cloning and world models on games such as Atari. This work seems like a rigorous starting point, however, impact may be limited by the small scale of the largest FLOPs budgets (<1e20) which remains orders of magnitude below the higher FLOP budgets in large-scale settings. For instance, Llama3 used 1e22 as the maximum budget for their scaling analysis.
Decision: Accept (poster)