Scaling Laws for Pre-training Agents and World Models

## Abstract

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 (e.g. between loss and optimal model size), also arise in world modeling and imitation learning.
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.

## 1 Introduction

Much progress in AI in the early 2020’s has been driven by increasing model size, dataset size, and training compute. Whilst conceptually simple, the importance of this practice has led to an emerging subfield studying the science 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.

The role of scale in pre-training is until now best understood in the context of large language models (LLMs).
Following the observation that the empirical relationship between loss and key scaling quantities can be accurately described by power laws [Kaplan et al., ], ensuing work studied the precise trade-off between model and dataset size [Hoffmann et al., ], as well as considerations about inference compute [Sardana and Frankle, ],
repeated training data [Muennighoff et al., ],
parameter counting [Pearce and Song, ], and more (Section 2).

In comparison, less is understood about scaling in embodied AI. Recent high-impact works suggest increasing model and dataset size can lead to ever more capable agents for two pre-training objectives; behavior cloning (BC) [Reed et al., , Baker et al., , Brohan et al., ] and world modeling [Hafner et al., , Hu et al., , Yang et al., , Bruce et al., ].
Such works typically demonstrate the benefit of scale through ablations over only a few model sizes, shown in terms of downstream agent performance, confirming the intuition that ‘bigger is better’ (Sartor and Thompson [] provide an aggregated analysis).
However, this leaves a large gap to the precise understanding of scale in LLMs, where for a given increase in compute, models can be sized optimally, and their expected performance accurately predicted.

This paper helps close this gap. Similar to the initial study of scale in LLMs, we focus on the effect of scaling on a generative pre-training loss (rather than on downstream agent performance, or reward- or representation-centric objectives), in 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 to several contributions, summarized in Figure 1.

- •
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 compression rate (number of tokens per observation) (Section 4.1, Figure 1a & b).
- •
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 1c).
- •
Scaling laws similar to those in LLMs can once again be observed in BC with one continuous encoding per observation (Section 4.2, Figure 1d).
- •
Our findings can be understood through small-scale language modeling experiments (Section 5).

Organization.
Section 2 provides detailed related work, contrasting the current understanding of scaling in embodied AI with other domains.
Section 3 introduces details for our main experiments, including the architectures considered and details of scaling laws analyses.
Section 4.1 & 4.2 present our main results in world modeling and BC respectively.
Section 5 presents insights behind our main results, including a set of tiny-scale language experiments mimicking aspects of our main experiments.
Section 6 discusses our findings and notes limitations in our work.

Figure: 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) However, moving to a single continuous embedding per observation remedies this (BC-CNN), producing prototypical scaling laws and a more balanced optimal model size coefficient.
Refer to caption: x1.png

## 2 Related work

Scaling laws origin. The term scaling laws is used throughout the engineering and physical sciences to denote power law relationships between two quantities,
e.g. duration of a volcanic eruption and the probability of it continuing [Cannavò and Nunnari, ].
The name derives from the scale-invariant(^1^11For two variables $x$ & $y$, the power law $y=ax^{b}$ is invariant to scaling $x$ by a constant $c$. Formally: $a(cx)^{b}=c^{b}ax^{b}\implies c^{b}y=c^{b}ax^{b}\implies y=ax^{b}$.) property of power laws.
While early work suggested that power laws could be good empirical descriptors of important variables in deep learning [Hestness et al., , Rosenfeld et al., ], it was Kaplan et al. [] who provided a comprehensive study of power laws in transformer LLMs, and popularized the usage of scaling laws in this context.

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. [] conducted a precise analysis into the trade-off of model and dataset size, finding they should be increased in equal proportions. This conflicted with Kaplan et al.’s suggestion that model size should be prioritized – an incorrect conclusion that Pearce and Song [] showed largely arose from counting only non-embedding parameters.

Many other aspects of LLM scaling analyses are beginning to be refined. Su et al. [] revisited the methodology used to find scaling coefficients.
Hägele et al. [] found that multiple independent cosine schedules could be reproduced more efficiently through a constant learning rate with multiple short decays, or stochastic weight averaging.
Pearce and Song [] & Porian et al. [] found that well-tuned constant learning rates were sufficient to recover certain coefficients.
Bi et al. [] study the effect of various hyperparameters on scaling.
Muennighoff et al. [] looked at repeated epochs, finding up to four epochs produce negligible departures from the infinite data regime. Sardana and Frankle [] factored in inference compute to the definition of what is compute-optimal. Isik et al. [] 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, , Maloney et al., , Bahri et al., ].

Scaling in embodied AI. Compared to LLMs, our understanding of scale in embodied settings is less advanced. Early successes in competitive games showed that running reinforcement learning (RL) at scale could surpass human performance, e.g. [Silver et al., , Berner et al., ].
In self-play RL, power laws were observed between certain quantities by Neumann and Gros [].
Meanwhile, Hilton et al. [] noted that, in general, reward signals do not follow power laws, and defined a transformation of reward (intrinsic performance) that created self-consistent scaling laws.

Inspired by the effectiveness of scaling in LLMs, embodied AI research has recently begun to explore the effectiveness of generative pre-training objectives on offline datasets, when executed at scale.
This includes behavior cloning objectives in video games [Baker et al., , Raad et al., ], robotics [Brohan et al., , , Padalkar et al., , Bousmalis et al., ], or multiple domains [Reed et al., ], as well as world modeling objectives [Hu et al., , Yang et al., , Bruce et al., ].
In these studies, the benefit of scale is generally shown through increasing model size on a specific downstream task of interest (e.g. measured by completion rate) – an aggregated survey is provided by Sartor and Thompson [].

Whilst such studies provide strong evidence that scale can be effective in embodied domains, the complexity introduced by downstream task evaluation makes quantifying the effects, and asking more nuanced questions very challenging (e.g. how to trade-off model and dataset size).
To address this, we argue that a better starting point for studying scale in these generative pre-training approaches, is to study the effect of scale on pre-training loss. Whilst downstream task performance may be of ultimate interest, LLM research evidences the utility of focusing on this clean intermediate signal, which is more straightforward to analyze.
Few prior works have taken such an approach. In Appendix A we contrast our BC efforts with Tuyls et al. [] who provide a valuable initial study on the relationship between compute and pre-training loss for BC. We also discuss scaling analyses of image and video modeling, which are related to world modeling.

## 3 Methodology

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.1 Tasks

We consider trajectories constructed as sequences of alternating observations $\mathbf{o}_{t}$ and actions $\mathbf{a}_{t}$ for timestep $t\in\mathbb{N}$. In this work, observations are always images, $\mathbf{o}_{t}\in\mathbb{R}^{3\times w\times h}$ and any continuous actions are discretized during preprocessing leaving, $\mathbf{a}_{t}\in\{0,1\}^{d_{a}}$.

Given this data format, we consider two tasks. World modeling (WM) [Ha and Schmidhuber, ] predicts future observations 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, ]. Behavior cloning (BC) predicts the future actions that the dataset’s demonstrators take [Bakker et al., ]. This creates a policy that can be directly used to act in the environment, either as-is or following further fine-tuning.
Concretely, these two tasks require modeling the following quantities,

$$ World modeling: $\displaystyle\;\;P(\mathbf{o}_{t+1}|\mathbf{o}_{t}\dots\mathbf{o}_{t-k}, \mathbf{a}_{t}\dots\mathbf{a}_{t-k}),$ (1) Behavior cloning: $\displaystyle\;\;P(\mathbf{a}_{t}|\mathbf{o}_{t}\dots\mathbf{o}_{t-k},\mathbf{ a}_{t-1}\dots\mathbf{a}_{t-k-1}).$ (2) $$

This work focuses on generative pre-training aiming to model this full conditional probability distribution. We leave a study of scaling laws for alternative objectives, e.g., explicitly targeting representation learning [Nair et al., ] or reward-centric models [Hafner et al., ], to future work.

### 3.2 Architectures

Figure: Figure 2: 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.
Refer to caption: x5.png

All experiments revolve around GPT-2 style causal transformers [Radford et al., ] as the core of the model. However we consider two different methods for inputting image observations, summarized in Figure 2. Section 3.4 details how we measure the model size of each.

Tokenized architecture. Our first architecture tokenizes each image observation into multiple discrete tokens. This is done with a frozen VQGAN encoder $\operatorname{Enc}_{\theta}(\mathbf{o}_{t})\to\mathbf{z}_{t}$, where $\mathbf{z}_{t}\in\{1,2,...,V_{o}\}^{d_{z}}$, for vocabulary size $V_{o}$ and latent dimension $d_{z}$.
Discretized actions are mapped to a non-overlapping vocabulary.
Following tokenization, training sequences take the form,

$$ $[z_{t}^{1},z_{t}^{2},...,z_{t}^{d_{z}},a_{t}^{1},a_{t}^{2},...,a_{t}^{d_{a}},z _{t+1}^{1},z_{t+1}^{2},...z_{t+1}^{d_{z}},a_{t+1}^{1},a_{t+1}^{2},...,a_{t+1}^ {d_{a}},...],$ (3) $$

where each item of the sequence is an integer within our vocabulary. A transformer is then trained to maximize the likelihood of either the latent image tokens (world modeling), or action tokens (BC).

This tokenized architecture is widely used both in world modeling [Micheli et al., ] and BC tasks [Bousmalis et al., ].
Gato [Reed et al., ] used a similar design but with continuous patches rather than discrete tokens.
Our implementation tests both a ‘small’ (28M parameters, $d_{z}=256$) and ‘large’ (150M parameters, $d_{z}=540$) VQGAN – further details in Appendix B.

CNN architecture.
Our second architecture differs in two ways. 1) Each image observation is input into the transformer 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 $P(\mathbf{a}_{t}|\dots)\approx\prod_{i=1}^{d_{a}}P(a_{t}^{i}|\dots)$. A single forward pass of the transformer is needed per action prediction.

This produces an architecture similar to Baker et al. [] (VPT additionally used a transformer-XL and a refined hierarchical action space).
Our implementation uses an Impala-style [Espeholt et al., ] CNN with 0.6M parameters for embedding image observations.

### 3.3 Datasets

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.

- 1.
Dataset size. Repeated training on the same data alters the effect of scaling.
Therefore, datasets should be large enough that all model sizes use a low number of training epochs.
- 2.
Dataset diversity. Both the behavior and environment must contain enough richness and variety that pre-training loss does not saturate across the model sizes tested.

Many existing benchmark datasets fail to fulfill these criteria – if not due to limited size, then because behavior is generated from a pre-trained fixed policy, or the environment is too simple.

Our work instead focuses on a dataset of human behavior collected in a video game named Bleeding Edge.
This is a fast-paced 4 vs 4 multiplayer game, with a range of characters, abilities and maps.
Game play is highly complex due to the cooperative and competitive dynamics. Success requires selecting high-level strategies (e.g. choosing which map regions to fight for), as well as fine-grained reactive control during combat.
Figure 9 shows example sequences from our dataset.

Supported by the game’s developer Ninja Theory, we compiled a dataset of 8.6 years of anonymized game play, containing both image observations and controller actions. We refer to this as the 7 map dataset. We also use a subset of this for some experiments, of around 1.1 years from a single map, which we name the Sky Garden dataset. Appendix B.3 provides further technical details.

### 3.4 Scaling analysis methodology

We are interested in studying the relationship between several quantities defined below.

- •
Model size $N$, the total number of trainable parameters (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 and Song [].
- •
Dataset size $D$, the total number of inputs the transformer sees during training. For WM-Token and BC-Token this is $d_{z}+d_{a}$ per observation & action pair, and for BC-CNN this is one per observation & action pair.
- •
Compute $C$, the number of floating point operations (FLOPs) used during training. The common approximation of $C=6ND$ [Kaplan et al., ] is used.
- •
Loss $L$, the standard classification cross-entropy loss (all targets are discretized).
We assume training loss is an accurate proxy for test loss (Appendix B.3.1 analyzes further).

More specifically, we are interested in ‘compute-optimal’ versions of each quantity. For loss, this is defined as the minimal loss possible for a given FLOPs budget,

$$ $\displaystyle L_{\text{optimal}}(C)$ $\displaystyle=\underset{\text{s.t.}\;C=6ND}{\operatorname{min}}L(N,D),$ (4) $$

where $L(N,D)$ is the empirical loss achieved with an $N$ parameter model trained on $D$ tokens. We further define optimal model and dataset sizes as the configuration that produce this minimal loss given a FLOPs budget,

$$ $\displaystyle N_{\text{optimal}}(C),D_{\text{optimal}}(C)=\underset{N,D\;\text {s.t.}\;C=6ND}{\operatorname{argmin}}L(N,D).$ (5) $$

Scaling analysis.
The heart of scaling law analysis is fitting power law relationships predicting these compute-optimal quantities. For predicting optimal model and dataset size, we use,

$$ $\displaystyle\hat{N}_{\text{optimal}}(C)$ $\displaystyle=a_{0}C^{-a}\;\;\;\;\;\;\hat{D}_{\text{optimal}}(C)=b_{0}C^{-b},$ (6) $$

with fitted constants $a_{0},a,b_{0},b$. We consider two methods to fit these relationships, introduced by Hoffmann et al. []. Their Method 1, which we term Frontier fit, classifies efficient models as those falling on the efficient frontier (see Figure 1). Coefficients can then be estimated straightforwardly through a line of best fit on a plot of FLOPs vs parameters or data for these efficient models.

Frontier fit is our preferred method when available – it avoids making any assumptions about the training curves, directly fitting the best models observed.
However, it requires training models past the point where they are the optimal configuration (seen on a loss-FLOPs plot as overlapping 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. [], which we term Parametric fit. This fits the coefficients $\alpha,\beta,N_{c},D_{c},E$ to a parametric loss form,

$$ $\displaystyle\hat{L}(N,D)$ $\displaystyle=\frac{N_{c}}{N^{\alpha}}+\frac{D_{c}}{D^{\beta}}+E,$ (7) $$

to the empirical training curves.
In our implementation, we use SciPy’s curve_fit function. We then find $a=\beta/(\alpha+\beta),b=\alpha/(\alpha+\beta)$. 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 and Song [],

$$ $\displaystyle\hat{L}_{\text{optimal}}(N,D)$ $\displaystyle=c_{0}C^{-c}+E.$ (8) $$

We again use the curve_fit function, fitted to models along the efficient frontier. During fitting, we set bounds on the variables: $c_{0}\in[0,\infty],c\in[-1,1],E\in[0.1,\infty]$.

Training details. While early scaling studies conducted sweeps over multiple cosine decays of differing lengths [Kaplan et al., , Hoffmann et al., ], follow up work found this redundant [Pearce and Song, , Hägele et al., , Porian et al., ]. We follow the approach of using a constant learning rate per model, so each requires only one training run.
We aim to 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 B.

## 4 Scaling analysis in embodied AI

This section presents our main results.
We begin by considering 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.

**Table 1: Summary of fitted scaling coefficients for our main experiments. Note that we favor the Frontier fit when available, and only use the Parametric fit for BC-Token-540 (see Section 3.4).**
|  | Frontier fit | Parametric fit |  |  |
| --- | --- | --- | --- | --- |
| Experiment | $N_{\text{optimal}}\propto C^{a}$ | $D_{\text{optimal}}\propto C^{b}$ | $N_{\text{optimal}}\propto C^{a}$ | $D_{\text{optimal}}\propto C^{b}$ |
| WM-Token-256 | 0.49 | 0.51 | 0.52 | 0.48 |
| WM-Token-540 | 0.62 | 0.37 | 0.78 | 0.22 |
| BC-Token-540 | N/A | N/A | 0.32 | 0.68 |
| BC-CNN | 0.66 | 0.34 | 0.47 | 0.53 |

### 4.1 Scaling analysis in world modeling

Figure: Figure 3: WM-Token scaling with $d_{z}=$256 tokens per image observation. Left shows the parametric fit. Middle & right show the frontier fit estimating optimal model & dataset size respectively.
Refer to caption: x6.png

Figure: Figure 4: WM-Token scaling with $d_{z}=$540 tokens per image observation. Left shows the parametric fit. Middle & right show the frontier fit estimating optimal model & dataset size respectively. Compared to the results for WM-Token-256, the power law coefficient for $N_{\text{optimal}}$ increases from $0.49$ to $0.62$.
Refer to caption: x9.png

Figures 3 & 4 present our results for the task of world modeling, with the scaling law coefficients summarised in Table 1.
For WM-Token-256 we find that the optimal coefficients for model and dataset size are both $\approx 0.5$, e.g. one should increase both model and dataset size in the same proportions.
This matches the scaling laws observed in LLMs [Hoffmann et al., ].
Increasing the number of tokens per image to $540$ for WM-Token-540 changes the optimal trade-off between model and dataset size, skewing towards model size; $N_{\text{optimal}}=0.62$, $D_{\text{optimal}}=0.37$. We discuss this further in Section 5.3.

### 4.2 Scaling analysis in behavior cloning

Figure: Figure 5: BC-Token scaling with $d_{z}=$ 540 tokens per image observation. Models above 2M parameters do not saturate over the FLOPs range considered and coefficients can not be inferred using the frontier fit method.
Refer to caption: x12.png

Figure: Figure 6: BC-CNN scaling. Left shows the parametric fit. Middle & right show the frontier fit estimating optimal model & dataset size respectively. Compared to the results for BC-Token the model sizes considered compute-optimal are considerably larger. The power law coefficient for $N_{\text{optimal}}$ also increases significantly from $0.32$ to $0.66$ skewing towards scaling model size as opposed to dataset size when scaling up compute.
Refer to caption: x15.png

We present our results on the scaling law coefficients for BC-Token in Figure 5.
Despite sharing an architecture with WM-Token-540 we now observe the opposite dependence on model and dataset sizes.
The coefficients skew heavily towards dataset size; $N_{\text{optimal}}=0.32$, $D_{\text{optimal}}=0.68$ (compared to $N_{\text{optimal}}=0.62$, $D_{\text{optimal}}=0.37$ – explained in Section 5.1).
Furthermore, under the same compute budget the compute-optimal model sizes are significantly smaller.
For a compute budget of $10^{18}$ and $10^{19}$ FLOPs we find that model sizes of $2$M and $11$M are compute-optimal for BC-Token-540 compared to $27$M and $110M$ for WM-Token-540.
In our experiments, we observe the losses for the BC-Token models take much longer to plateau leading to less overlap between model sizes.
This results in the frontier fit not being suitable for accurately estimating the scaling law coefficients, hence we rely on the parametric fit for these results.

To better understand the change in the scaling law coefficients, we now consider the BC-CNN architecture for the task of BC in Figure 6.
For this architecture, we observe that the coefficients now skew towards model size (similarly to those in [Tuyls et al., ]), with $N_{\text{optimal}}=0.66$, and $D_{\text{optimal}}=0.34$.
Section 5.2 provides more intuition on the differences between the WM-Token and BC-Token setups that lead to this change.

Further to studying the differences in scaling law coefficients between tasks and architectures, we also study the accuracy of extrapolation.

### 4.3 Extrapolation in world modeling

Figure: Figure 7: Testing the extrapolation capability of our derived scaling law for WM-Token-256 by training an $894$M 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 of training tokens (dotted lines) and our actual training run.
Refer to caption: x18.png

To test the extrapolation accuracy of our derived scaling laws, we train a $894M$ parameter WM-Token-256 model with an order of magnitude more compute than used for the scaling law analyses.
Figure 7 presents both the learning curve as well as the extrapolated lines derived from the Frontier fit method.
We take the point with the loss value closest to our extrapolated loss curve ($\sim 1.58\times 10^{21}$FLOPS), 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, demonstrating 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 $894$M parameter model, which was not extensively tuned due to compute requirements.

## 5 Further analysis

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 identical architecture and dataset? (Section 5.1)
- •
Q2: Why does moving from BC-Token to BC-CNN resolve this issue? (Section 5.2)
- •
Q3: Why does increasing the amount of tokens per image observation (256 to 540) lead to an increase in the optimal model size coefficient (0.49 to 0.62)? (Section 5.3)

### 5.1 Q1: BC-Token vs. WM-Token

The lack of saturation of BC-Token models compared to WM-Token models can be attributed to two factors.
The first is a sparser loss. A single observation-action pair is discretized into $d_{z}+d_{a}$ total tokens. With the large VQGAN tokenizer, world modeling receives supervision for $d_{z}/(d_{z}+d_{a})=540/556\approx 97\%$ tokens, while BC is supervised for $d_{a}/(d_{z}+d_{a})=16/556\approx 3\%$ of tokens.

The second factor is the granularity of the targets. The large tokenizer creates a world modeling vocabulary size of 4096. Each vocabulary item roughly corresponds to a specific 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, a player can take the same action in multiple different situations – continue straight could be used to escape an 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.

To demonstrate the effect of these two factors on optimal model size coefficients, we run a set of tiny-scale experiments in language modeling.
Transformers are trained on next-character prediction, on a dataset of Shakespeare text(^2^22Shakespeare character dataset from: [https://github.com/karpathy/nanoGPT](https://github.com/karpathy/nanoGPT)) using a single character for each token. Model sizes are varied from 4k parameters to 17M parameters. Context length is fixed at 16 characters/tokens.

Figure 8 (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 a 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-Token. We compute optimal 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 to 0.32, indicating that the proposed mechanisms explain our findings.

Figure: Figure 8: Training curves and parametric fit for character modeling experiments. The standard dense LLM loss has been modified to reflect properties of BC – a sparse loss (one of 16 tokens), and then additionally super-classing the targets into two classes.
Refer to caption: x21.png

### 5.2 Q2: BC-Token vs. BC-CNN

Despite the same non-granular loss signal, why does switching architecture from BC-Token to BC-CNN makes the loss of similar model sizes plateau under a much smaller compute budget?

Consider each architecture using a transformer with 1M parameters.
Observe from Figure 2 that BC-Token receives $d_{z}+d_{a}=556$ inputs for every action $\hat{\mathbf{a}}_{t}$ it predicts, while BC-CNN receives just one input for every action predicted. Hence, BC-Token uses around 556 times more compute in its action prediction ($556\times 2\times 1M\approx 1\times 10^{9}$ FLOPs) than BC-CNN ($1\times 2\times 1M\approx 2\times 10^{6}$ 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.

### 5.3 Q3: WM-Token-256 vs. WM-Token-540

Finally, we seek to understand why the optimal model size coefficient increases when moving from the 256 to the 540 token VQGAN.
As the number of tokens per image observation are increased, the compression rate of the tokenized representation decreases.
We would expect that each individual 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 of examples would need to 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.

For further insight into how compression affects the optimal model size coefficient, we ran a small scale experiment in language modeling using two text representations; 1) ASCII character-level tokenization. (low compression) 2) GPT-2 tokenizer (high compression). We used the BookCorpus dataset [Zhu et al., ], and trained models past their compute-optimal point, so the Frontier fit method could be used for coefficient estimation.

Appendix C shows results. Under the character-level tokenizer (low compression), we find $N_{\text{optimal}}\propto C^{0.66}$. For the GPT-2 tokenizer (high compression), we find $N_{\text{optimal}}\propto C^{0.44}$. The key takeaway is that the more compressed representation has led to a lower optimal model size coefficient, which reflects what we see in the world modeling experiments.

## 6 Discussion & conclusion

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 recover scaling laws similar to those established in the LLM literature. Establishing such a link is key to making efficient use of available resources, and to training compute-optimal models.

Considering the task of world modeling, we find that models can be smoothly scaled following best practices and insights from the LLM literature. Surprisingly, the scaling coefficients for our WM-Token-256 architecture very closely match those established for LLMs. Comparing to our WM-Token-540 model and additional analysis, we further establish that scaling is affected by the tokenizer’s compression rate.

Turning to pre-training BC policies for agents, the choice of architecture is extremely important in determining optimal scaling behavior.
When using architectures with tokenized image observations, dataset size should be increased much more rapidly than model size. Meanwhile, for BC-CNN architectures, model size should be increased faster than dataset size.

Limitations.
While we show that scaling laws can be precisely described in the infinite data regime and for appropriate architectures, future work is needed to establish scaling laws for alternative models and under varying dataset quality. In addition, we focus on loss as an intermediate quantity that can be effectively optimized in pre-training. Many additional considerations are required for effective AI models, such as downstream task performance and model inference times. How valuable scaling laws can be in providing insights relevant to those choices remains an open question.

#### Ethics Statement

Data for this project was provided via a partnership with *Ninja Theory*, who collected a large corpus of human gameplay data for their game *Bleeding Edge*. Data collection was covered by an End User License Agreement (EULA) and our use of the data was governed by a data sharing agreement with the game studio, and approved by our institution’s IRB. This data was recorded between September 2020 and October 2022. To minimize risk to human subjects, any personally identifiable information (Xbox user ID) was removed from the data. The resulting data was cleaned to remove errors and data from bad actors.

## Appendix A Extended related work

Here we contrast more granular details of several related works.

Tuyls et al. [] provide (amongst other things) a valuable initial study on the relationship between compute and pre-training loss for BC. They focus on Atari and Nethack games, where the dataset is generated by a fixed pre-trained agent.
They test one architecture – a single-layer LSTM network, scaling width-wise to increase parameter count. They find the model size coefficient $N_{\text{optimal}}\propto C^{a}$ varies between 0.58 to 0.79 for different games.

By contrast, our work focuses on transformer-based models, and we consider two variants of architecture. We also train on human behavior, rather than pre-trained agents (which could be viewed as model distillation).
Additionally, we go beyond BC and also study scale in world modeling.

Scaling laws in other domains.
Scaling laws have also been observed in auto-regressive modeling of other modalities. Most relevant to our world modeling task are those of video and images. Henighan et al. [] found the optimal trade off between model and dataset size to match their reported LLM coefficient ($N_{\text{optimal}}\propto C^{0.7}$) and was not affected by tokenizer. Our experiments offer different findings in the domain of world modeling on human data – we use updated methodologies to measure this trade-off [Hoffmann et al., ], finding it is affected by the tokenizer. Tian et al. [] show scaling laws emerge in image modeling with non raster-ordered next-token prediction.

## Appendix B Scaling experiments further details

### B.1 Hyperparameters

We trained three VQGANs from scratch with reconstruction losses.

- •
BE-Small. Based on Esser et al. [], uses $d_{z}=256$, $V_{o}=4096$, $h=w=128$, with 28M parameters, and a CNN design. It was trained on a single SkyGarden Bleeding Edge map.
- •
BE-Large. Based on Yu et al. [], uses $d_{z}=540$, $V_{o}=4096$, $h=180$, $w=300$, with 150M parameters, and a vision transformer design. It was trained on all seven Bleeding Edge maps.

BC-CNN details. We use $h=w=128$. The $0.6M$ paramter CNN is similar to that used by [Baker et al., ], however it uses ConvNext blocks [Liu et al., ].
The CNN produces an embedding of size $1024$ which is then put through a linear layer to obtain a vector matching the transformer’s embedding dimension.

Transformer configurations are given in Table 2. We describe the parameters for the WM-Token architecture.
Note that MLP layers are four times the width of embed dim.

**Table 2: Transformer configurations. Here $N$ is listed for the tokenized architectures. Parameter count varies slightly for BC-CNN due to 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 |
| 11M | 6 | 6 | 360 |
| 15M | 4 | 4 | 512 |
| 27M | 8 | 8 | 512 |
| 52M | 10 | 10 | 640 |
| 110M | 15 | 12 | 768 |
| 206M | 16 | 16 | 1024 |
| 894M | 23 | 14 | 1792 |

### B.2 Training details

All transformers are trained with a variant of nanoGPT [Karpathy, ] using PyTorch Lightning [Falcon and The PyTorch Lightning team, ].

**Table 3: Hyperparameters for WM-Token with $d_{z}=$256 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_{z}=$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 5: Hyperparameters for BC-Token with $d_{z}=$540 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 6: 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 | 10 | 2560 | 0.0003 |
| 26M | 10 | 10 | 2560 | 0.0003 |
| 51M | 10 | 10 | 2560 | 0.0003 |

### B.3 Dataset details

Figure: Figure 9: Example trajectories from a dataset of 8.6 years of human gameplay in the video game Bleeding Edge across 7 maps.
Refer to caption: extracted/5978876/01_images/7map_frames.png

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 7 Maps dataset comprised 60,986 matches, yielding 530,713 individual player trajectories (each around 9 minutes), totaling 27.89 TiB on disk. This amounted to around 8.6 years of gameplay. After downsampling to 10Hz (the frequency models are trained on), this equated to 1.63B frames. This was then divided into training / validation / test sets by dividing the matches with an 80: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 x and y values of the left and right joysticks into eleven buckets.

#### B.3.1 On the infinite data regime

We wish to 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 1 alongside these numbers confirms that models remain in the infinite data regime for all our experiments.

WM-Token-540, BC-Token-540.
We trained on the 7 maps dataset, with 1.63B observation-action pairs. Models used the tokenized architecture with the large VQGAN, so each observation-action pair creates $540+16=556$ transformer inputs, for a total of 1.63B$\times 556=906$B training tokens.
Muennighoff et al. [] observe that tokens may be reused up to four times with negligible departure from the infinite data regime. This produces 3.6T tokens.
For a 200M parameter model the compute allowed by the infinite data regime is $C=6ND=6\times 200\text{M}\times 3.6\text{T}=4.3\times 10^{21}$ FLOPs.

WM-Token-256. 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$\times 4=386$B effective tokens.
For a 200M parameter model the compute allowed by the ‘infinite data regime’ is $C=6ND=6\times 200\text{M}\times 386\text{B}=4.6\times 10^{20}$ FLOPs.

BC-CNN. Trained on 7 maps dataset, but now with one token per observation-action pair, this creates a possible $1.63\text{B}\times 4=6.52\text{B}$ effective tokens. A 50M parameter model uses $C=6ND=6\times 50\text{M}\times 6.52\text{B}=2.0\times 10^{18}$ FLOPs.

## Appendix C Further analysis details

Experimental results supporting Section 5.3.

Figure: Figure 10: Relating to Section 5.3, character-level (low compression). Utilising the frontier fit (middle and right) we derive the power law coefficient for $N_{\text{optimal}}$ as $0.66$.
Refer to caption: x24.png

Figure: Figure 11: Relating to Section 5.3, GPT-2 tokenizer (high compression). Utilising the frontier fit (middle and right) we derive the power law coefficient for $N_{\text{optimal}}$ as $0.44$, an increase from $0.66$ in Figure 10 found when utilising a lower compression character-level tokenizer.
Refer to caption: x27.png

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: ## Questions
1. Would incorporating large-scale real-world robotics data help validate the paper's scaling law claims beyond simulation environments? The current findings, while valuable, are limited to the simulation world.
2. Could the authors demonstrate the practical value of lower pre-training losses by evaluating a subset of pre-trained models on downstream tasks? This would help establish a clear correlation between pre-training performance and downstream task effectiveness.
soundness: 2
strengths: ## Strengths
1. **Visual Encoder Architecture Exploration:** The study explores scaling laws across various visual encoding architectures and pre-training tasks including BC and WM, showing that model scaling behavior in embodied AI resembles language modeling. For example, WM tasks with 256-token representations require scaling model and dataset sizes equally $N(optimal) ∝ C^{0.49}$ and $ D(optimal) ∝ C^{0.51} $.
2. **Tokenizer compression:** The authors find that tokenizer compression rates significantly affect scaling laws. For instance, increasing token compression from 256 to 540 tokens per image in the WM task shifts optimal model scaling to $N(optimal) ∝ C^{0·62}$, favoring larger model sizes.
summary: ## Summary
This paper investigates the impact of model size, dataset size, and compute on embodied AI tasks like behavior cloning (BC) and world modeling (WM). Using transformers with tokenized inputs or CNN embeddings, the authors explore scaling behavior using 8.6 years of gameplay data from *Bleeding Edge*. They find that scaling laws—typically observed in language modeling—also apply to embodied AI tasks, with different trade-offs influenced by task, architecture, and tokenization. For instance, in the WM task with 540-token representations, optimal model size scales as $N(optimal) ∝ C^{0.62}$, whereas BC tasks see optimal model size scaling as $ N(optimal) ∝ C^{0.32} $ with the same tokenized setup.

## Quality & Clarity
The paper is written clearly, with effective visualizations such as Figures 3 and 4 showing the loss curves and power law fits for scaling laws across various model sizes. The methodology in Table 1 provides fitted coefficients across model-dataset configurations. In Section 5, several design choices, such as the comparison of architectural choices for BC-Token versus BC-CNN, are explored, and explanations are provided for why specific scaling trends arise in each architecture.

## Recommendation
The paper contributes to the understanding of scaling laws for embodied AI, particularly with evidence of model-dataset trade-offs for BC and world modeling tasks on a single dataset. However, the contributions are limited by their similarity to previous work (e.g., Tuyls et al.) and the relatively narrow focus on architectural adjustments rather than task diversity or downstream metrics. I currently recommend rejection but would be willing to change my stance if the aforementioned weaknesses are addressed.
weaknesses: ## Weaknesses
1. **Lack of real-world data:** The scaling insights are based on a single simulation task. Previous studies have demonstrated that simulation performance often does not correlate well with real-world robot performance [1]. The analysis would benefit from validation on openly available real-world datasets such as OpenX[2] and DROID[3].
2. **Downstream task performance:** The paper focuses exclusively on scaling pre-training loss without examining generalization or performance on downstream tasks, which is crucial for embodied AI applications. Including evaluations on even a select set of downstream tasks would significantly strengthen the practical relevance of the findings.

## References
1. An Unbiased Look at Datasets for Visuo-Motor Pre-Training: Dasari et al.
2. Open X-Embodiment: Robotic learning datasets and RT-X models: OXE Collaboration
3. DROID: A Large-Scale In-the-Wild Robot Manipulation Dataset: Khazatsky et al.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: Please refer to the weakness section.
soundness: 3
strengths: - The paper focuses on an important problem that will have a large impact on the community.
summary: This paper explores the scaling law within an embodied agent setting, analyzing two tasks: world modeling and behavior cloning. The results indicate that a scaling law exists for both tasks when additional compute resources are provided.
weaknesses: Although it is a good attempt at analyzing the scaling law, the analysis lacks comprehensiveness. 
 - The relationship between loss and dataset size, which is mentioned in the introduction, is not presented. For example, in line 53, the authors state, “The optimal trade-off between model and dataset size in world modeling is influenced by the tokenizer’s compression rate (number of tokens per observation) (Section 4.1, Figure 1a & b).” However, Figures 1a and 1b primarily illustrate the effects of computation with different tokenizers and model sizes, without clarifying the influence of dataset size on the results.
 - The experiments utilize a frozen VQGAN visual encoder, which implies that the final claims regarding tokenizers are inherently limited by the performance of the frozen encoder.
 - The architectural analysis includes a set of BC-CNN experiments under the BC loss, but it is unclear why the CNN experiment was not conducted within the world modeling tasks.

---

confidence: 4
contribution: 1
initial_rating: 3
presentation: 3
questions: See above

What is this paper doing or concluding that is above prior work?

Why should I trust this proxy loss (train versus test and not final performance) over other works which directly measure the variable of interest?
soundness: 3
strengths: The paper is well presented and generally clear.

The analyses are well explained and clear to read.

There are a lot of interesting analyses done, the authors really tried a lot of things here and have a lot of results to pick through and analyze.
summary: This paper studies scaling laws in the context of agents (specifically in a video game) to study whether and how these laws exist in this setting. The paper looks at how dataset and model size can be optimally scaled to optimize loss for BC and world modeling (next state prediction). The authors find a power law relationship and show how things such as tokenization can effect the optimal scaling parameters.
weaknesses: The premise of what this paper is doing is kind of questionable. Specifically, the thing that is being measured is training (will get to that) loss on a BC or state prediction objective, rather than a direct measure of downstream performance. As the authors point out in their related work, most other papers such as Hu et al 2023 do measure downstream performance. The author's justification for this choice then is that this adds complexity and it might be difficulty to ask "more nuanced questions" such as how to trade-off model and dataset size. But this isn't really further justified any further. Why couldn't we do a study looking at that tradeoff on the downstream task. More importantly, this paper draws conclusions based on the training loss for a *proxy* to what researchers actually care about (how well the final agent actually does at test time), and given the choice between looking at this paper's measurements and one that evaluates this, I don't know why you would prefer this. There isn't really even much analysis about why this is a good proxy metric (like some kind of finding that things with lower loss in this paper do in fact do better on the final task, even if the curves are not as smooth)

Not only this, but the paper isn't even looking at the Test loss in these analyses, it's looking at train loss. From S3.4 "We assume training loss is an accurate proxy for test loss. (Appendix B.3.1 analyzes further)." I find this assumption both baffling and completely unjustified in the text of the paper. The cited appendix merely looks at the size of the data they trained on and says "The compute allowed by the infinite data regime is 6ND" (and so they say, they are above that. There isn't a citation for this anywhere, and I am not familiar enough with the literature to tell if this is a thing other works have done. The original scaling law paper as far as I can tell does not do this. But I still do not understand the decision. If you have this large a dataset, surely you have enough data to have a held-out set? Again, scaling laws are not something I was intimately familiar with, so if there is a good justification for this, someone please let me know. But this paper does not justify it and on it's face "assume train and test are the same" feels like an ML sin.

The results are also only validated on a single environment / datasets. This seems like it really is limiting the potential applicability of the paper. What if the conclusions or scaling parameters found in this paper are idiosyncratic to this particular game?

I say this somewhat advisedly, but I am not actually sure the models studied here are large enough to fully justify all the conclusions. Most of the experiments use 206M parameters with many having even smaller (e.g. 27M for Figure 5). The only experiment with a larger size is Figure 7 at 894M. This isn't my biggest criticism, and I know compute is quite expensive, but it feels like this is a major missing thing here.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: This paper makes promising contributions to the understanding of scaling laws in the context of pre-training for embodied AI, with clear strengths in empirical validation and methodological rigor. Addressing the identified gaps in hyperparameter transparency and broadening the scope of environments and architectures considered could further enhance its impact. Moreover, clarifying how this research fits into the broader context of embodied AI or RL, specifically the types of tasks and applications targeted and the importance of scaling laws for these areas would make the findings more relevant and engaging. Currently, the title is a bit ambiguous or too broad to understand the scope. I am happy to increase the score post-author response.
soundness: 3
strengths: #### Strengths:
1. **Empirical Validation:**
   - The paper provides evidence that scaling laws known from LLMs similarly apply to world modeling and BC, validated across some architectures and tokenization strategies.
   - Thoroughness of many results.
   - The systematic experimentation using different model sizes and computational scales support the credibility of the findings.

2. **Take on Tokenization and Architecture:**
   -  Impact of tokenization granularity on optimal model scaling is particularly novel, offering insights into how token compression influences model efficiency and effectiveness.
   - The comparison between tokenized and CNN based approaches in BC provides nice distinctions in how these architectures manage to scale.

3. **Methodology and Analysis:**
   - The paper details the methodologies employed, from the scaling analysis approach to breakdown of the computational and parameter scaling. This thoroughness can ensure clarity and reproducibility of the results.
   - The use of frontier fit and parametric fit methods to analyze scaling laws provides a usable framework for understanding the optimal trade-offs between model size and computational expense.
summary: Less is understood about scaling in embodied learning, but recent works suggest that increasing model and dataset size can lead to capable agents. This paper explores this hypothesis further by delving into scaling laws applicable to pre-training embodied agents and their world models, leveraging larger datasets and enhanced computational resources to improve performance. Focusing on both world modeling and behavior cloning, the study extends the understanding of how scaling affects these models beyond simple increases in parameter count, introducing dependencies on factors like tokenization and model architecture.
weaknesses: #### Areas for Improvement:
1. **Hyperparameter Selection and Rationale:**
   - While the paper provides a framework for understanding scaling laws, a more granular justification of the choices behind specific hyperparameters, particularly in the synthetic data generation and scaling analyses, would enhance its contribution. Specifically, selection of number of tokens per image in the tokenization process and the learning rates used across different model scales seem crucial to the study’s outcomes but lack a detailed rationale. Providing insights into how these parameters were optimized, may be through sensitivity analyses or referencing empirical benchmarks, could help.
   - Insights into the selection process for training configurations across different scales could further enhance the paper’s utility to practitioners.

2. **Broader Applicability and Generalization:**
   - The study provides insights into scaling laws within controlled experimental setups. However, to enhance the relevance and applicability of these findings to real-world applications, it would be beneficial to extend the analysis beyond the current datasets. Specifically, exploring how these scaling laws hold across datasets with higher environmental variability and some degree of realism or some evidence of extensibility.
   - Extending the analysis to include more diverse datasets or environmental complexities could help validate the laws observed.

3. **Limitations and Future Work:**
   - The paper acknowledges limitations in the scope of architecture and tokenization diversity. 

4. **Less emphasis on embodied AI:**
   - my main criticism concerns the paper's lack of a cohesive narrative explaining the relevance of its findings to embodied intelligence. It is currently less sound in terms of the necessity of the study, or why BC or world modelling techniques need this study particularly. is it because they are the most promising approaches.. or are they the most compelling representative candidates of model-free and model-based approaches or where are we heading with this analysis? Perhaps this can be easily fixed by revising some of the scaling law papers for sing-agent or multi-agent papers cited in the paper. Lastly, might be missing some relevant references and citations in the area, some are highlighted below and good to cite.

[1] GenRL: Multimodal-foundation world models for generalization in embodied agents, Mazzaglia et al

[2] Daydreamer: World models for physical robot learning, Wu et al

[3] ArCHer: Training Language Model Agents via Hierarchical Multi-Turn RL Zhou et al.

[4]Choreographer: Learning and Adapting Skills in Imagination, mazzaglia et al

Meta-review: This paper investigates scaling laws in embodied AI, focusing on behavior cloning (BC) and world modeling (WM) tasks. Using a large dataset from gameplay in Bleeding Edge, the study examines the interplay of model size, dataset size, and compute on pre-training losses. The authors demonstrates that scaling laws, similar to those found in language modeling, also apply to embodied AI tasks, revealing distinct trade-offs influenced by tokenization, task type, and architecture. Key findings include how token compression impacts the optimal balance between model and dataset size in WM tasks, and the difficulty of observing scaling laws in BC tasks under modest compute budgets. The aim is to infer insights into optimizng training setups and extending our understanding of scaling in embodied learning tasks.

Reviewers appreciated the novel exploration of scaling laws in embodied AI, noting scaling laws known from LLMs similarly applies to world modeling and BC. They highlighted the impact of tokenization compression on scaling, showing how increasing tokens per image shifts optimal scaling in WM tasks to favor larger models. The paper’s thorough analyses, systematic experimentation, and clear methodologies, including frontier and parametric fits, was also commended for supporting credibility and offering a framework to study trade-offs between model size and computational resources.

However, the reviewers raised concerns about the comprehensiveness of the evaluations for the paper’s claims. They noted the study’s reliance on a single dataset ("Bleeding Edge") and the lack of validation on real-world datasets, which limits its generalizability. Reviewers also questioned the use of pre-training loss as a proxy for downstream performance and called for direct evaluations on practical embodied AI tasks. Concerns were also raised about the relatively smaller-scale models used in experiments and insufficient justification for key hyperparametrs. In the rebuttal, the authors added world modeling results on RT-1 and conducted a metaanalysis on the pretraining loss vis-a-vis downstream task performance, along with providing significant additional details in Appendix D and E. While these efforts partially addressed concerns, in post-rebuttal discussion, it was noted that a broader evaluation across tasks and datasets is still needed for the technical claims of the work. 

The AC concurs with the unanimous recommendation of the reviewers and finds that a fresh round of reviews for a more comprehensive draft would be necessary.

Decision: Reject