Scaling Laws for Pre-training Agents and World Models

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.

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.

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.

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.

1.2 Increased task diversity

Some reviewers wished to know whether the patterns we have observed in our video game dataset of human behavior would hold in new environments and datasets. Inspired by this, we include a new set of world modeling experiments on the robotics RT-1 dataset (Appendix D), finding that scaling laws again emerge in this set up. Furthermore, we provide further evidence for our claim that the tokenizers compression rate (number of tokens per observation) affects scaling coefficients.

We trained a set of VQVAEs with $V_o = 4096$ and $z_o \in [16, 36, 64, 100, 256]$, for 40,000 updates on batches of 128 (reconstructions here: https://ibb.co/JrgBkKW). We then train a range of WM-Token model sizes $N \in [0.08M, 0.2M, 0.28M, 0.54M, 0.99M]$ on each VQVAE (e.g. https://ibb.co/9WNFKBV), and measure scaling coefficients using the frontier fit method.

First, we observe that scaling laws similar to those in our main set of experiments emerge, with predictably decreasing losses and similar optimal parameter scaling coefficients.

Secondly, by measuring $a$ where $N_\text{optimal} \propto C^a$ for each VQVAE, repeated three times, we find an that the optimal parameter scaling coefficient increases with number of tokens per observation (https://ibb.co/HNrzkms). This supports our claim that a higher compression rate leads to a lower optimal parameter coefficient.

1.3 Comparisons with prior work

Tuyls et al. [1] provide a valuable complimentary analysis to our own work. Here we emphasize several differences.

• 1. Half of our contributions focus on scaling laws for world modeling, Tuyls et al. conducted no investigation of this.
• 2. For our BC contributions, Tuyls et al. focused on scaling the width of CNN and LSTM models on datasets created by fixed policies. We consider transformer architectures on datasets of human behavior. We focus on understanding how architectural choices effect scaling coefficients -- surprisingly finding that they heavily influence optimal scaling strategy.
• 3. Section 8 of Tuyls et al. directly mentions that conducting scaling analyses on human datasets is an important direction for future work -- human datasets come with a broad distribution coverage and variation in skill level. This is exactly what our work focuses on.

References

[1] Scaling Laws for Imitation Learning in Single-Agent Games
[2] Hyperparameter Selection for Imitation Learning
[3] Learning Interactive Real-World Simulators

Thanks to all reviewers for the valuable feedback. Here we address points that were common to two or more reviewers. We respond to points raised by individuals in separate responses.

1.1 Pre-training loss vs. downstream task performance

Several reviewers queried the core premise of our work -- that pre-training loss is a useful thing to analyze and optimize for. Instead they requested a focus on online performance.

We begin by presenting two new pieces of evidence showing that pre-training loss is closely related to model quality. These have been incorporated into the paper (currently Appendix E).

• 1. World model evidence. We use metrics commonly used to assess the quality of the world models, e.g. [3], originally developed in the video generation literature. Conditioned on an initial real frame and a sequence of real actions, we compare the observations generated by the world model, with the real sequence of observations, measuring FVD, LPIPS and PSNR. Specifically, we generate 1024 videos of 10 seconds each. We perform this for various checkpoints on each model size of our WM-Token-256. This allows a plot of the checkpoint pre-training loss vs world model quality metric to be assessed (e.g. https://ibb.co/gScQQ2S). We find correlations of 0.77, -0.57, 0.83 for LPIPS, PSNR, FVD respectively. This evidences the strong relationship between pre-training loss and world model quality.
• 2. BC evidence. We conducted a meta-analysis of [1] to quantify the strength of relationship between pre-training loss and online return. This is possible in their set up for two reasons. 1) They use a simple environment that could be run cheaply as a gym environment, so all checkpoints could be evaluated. 2) Their data generating policy was an expert rules-based policy, so BC necessarily leads to an expert agent. By taking data points from their Figure 5a and b, we created a plot of pre-training loss vs. online return. This comes out with a correlation coefficient of -0.97 (https://ibb.co/wL0C4kq). This provides strong evidence that in the infinite data regime, pre-training loss is a good proxy for online performance.

In general, we agree with reviewers that when online performance can be measured, this is preferable to pre-training loss. However, scaling law analyses require evaluating a large number of checkpoints over multiple training runs. This is not usually feasible for two reasons.

1. For complex environments and applications, such as real robotics or modern video games (our work), measuring online performance is costly. Testing a single BC policy (following finetuning) in the Bleeding Edge environment in our work requires around half a day. Testing the ability to plan using a world model is even more expensive.
2. When pre-training datasets contain a mixture of skill-levels (our work) it is not clear that pre-trained models are 'out-the-box' aligned with high-skill behavior. Hence the name pre-training stage -- models require further fine-tuning or steering to elicit the desired behaviors. This further increases the cost of assessing online performance.

Given these reasons making online performance broadly inaccessible, we now outline qualitative arguments for why pre-training loss is a valuable metric to focus on.

• 1. Pre-training loss has proven valuable in LLM research. LLMs face a similar problem to embodied AI -- one cares about metrics such as factual correctness, reasoning capability, and user engagement, not pre-training loss. Yet these are expensive to measure, and the LLM community has made much progress by focusing on the optimization of the clean intermediate signal of pre-training loss. Why should this not hold for the embodied AI community?
• 2. Next-token prediction intuition. To improve at next-token prediction in world modeling and BC, models intuitively must know more about the environment and behaviors. In BC, better predicting a human's next action requires a model to better understand things about the game objectives the human's are trying to complete, the skill level of individual trajectories, and a variety of alternative behaviors that humans may choose to perform. All these things should create a better pre-trained checkpoint for specialization to a particular downstream task.
• 3. Validation loss vs. infinite data. Previous work has sometimes reported a imperfect relationship between validation loss and performance, e.g. [2]. However, such studies have been performed on datasets of limited size, with held-out validation sets, finding that some amount of over-fitting might be beneficial. Our work is conducted in the infinite data regime, where over-fitting is not possible, so previous insights may not be applicable.

[1] Scaling Laws for Imitation Learning in Single-Agent Games
[2] Hyperparameter Selection for Imitation Learning

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.

Reject