Scaling Offline Model-Based RL via Jointly-Optimized World-Action Model Pretraining
Authors: Jie Cheng
Ruixi Qiao
YINGWEI MA
Binhua Li
Gang Xiong
Qinghai Miao
Yongbin Li
Yisheng Lv
Venue: ICLR 2025 Poster (2025)
Keywords: reinforcement learning
offline reinforcement learning
world model
Primary area: reinforcement learning
TL;DR: We scale offline model-based RL through a jointly-optimized world-action model pretrained across multiple games, which achieves sample-efficient transfer to novel games.

Abstract: A significant aspiration of offline reinforcement learning (RL) is to develop a generalist agent with high capabilities from large and heterogeneous datasets. However, prior approaches that scale offline RL either rely heavily on expert trajectories or struggle to generalize to diverse unseen tasks. Inspired by the excellent generalization of world model in conditional video generation, we explore the potential of image observation-based world model for scaling offline RL and enhancing generalization on novel tasks. In this paper, we introduce JOWA: Jointly-Optimized World-Action model, an offline model-based RL agent pretrained on multiple Atari games with 6 billion tokens data to learn general-purpose representation and decision-making ability. Our method jointly optimizes a world-action model through a shared transformer backbone, which stabilize temporal difference learning with large models during pretraining. Moreover, we propose a provably efficient and parallelizable planning algorithm to compensate for the Q-value estimation error and thus search out better policies. Experimental results indicate that our largest agent, with 150 million parameters, achieves 78.9% human-level performance on pretrained games using only 10% subsampled offline data, outperforming existing state-of-the-art large-scale offline RL baselines by 31.6% on averange. Furthermore, JOWA scales favorably with model capacity and can sample-efficiently transfer to novel games using only 5k offline fine-tuning data (approximately 4 trajectories) per game, demonstrating superior generalization.

Published as a conference paper at ICLR 2025
SCALING OFFLINE MODEL-BASED RL VIA JOINTLY- OPTIMIZED WORLD-ACTION MODEL PRETRAINING
Jie Cheng' 12, Ruixi Qiao'.2, Yingwei Ma², Binhua Li
Gang Xiong Qinghai Miao?, Yongbin Li Yisheng Lv
State Key Laboratory of Multimodal Artificial Intelligence Systems,
Institute of Automation, Chinese Academy of Sciences
Schooo of Artificial Intelligence, University of Chinese Academy of Sciences
Alibaba Group
ABSTRACT
A significant aspiration of offline reinforcement learning (RL) is to develop a generalist agent with high capabilities from large and heterogeneous datasets. However, prior approaches that scale offline RL either rely heavily on expert tra- jectories struggle to generalize to diverse unseen tasks. Inspired by the ex- cellent generalization of world model in conditional video generation, we ex- plore the potential of image Abservvtionnbasee world model for scaling offline RL and enhancing generalization on novel tasks. In this paper, we introduce World-Action model, offine model-based RL agent
JOWA: Cointly-Optiminn pretrained on multiple Atari games with billion tokens data to learn general- purpose representation and decision-making ability. Our method jointly opti- mizes a world-action model through a shared transformer backbone, which sta- bilize temporal difference learning with large models during pretraining. More- over, we propose a provably efficient and parallelizable planning algorithm to compensate for the Q-value estimation error and thus search out better policies. Experimental results indicate that our largest agent, with 155 million parame- ters, achieves 78.9% human-level performance on pretrained games using only 10% subsampled offline data, outperforming existing state-of-the-arl large-scale offline RL baselines by 71.4% on averange. Furthermore, JOWA scales favor- ably with model capacity and can only 5k offline fine-tuning data (approximately 4 trajectories) per game, demon- strating superior generalization. The code and checkpoints will be released https:: htt///gttubb... com com//CReinfooreee...coooooooooe.
transfer to novel games using
INTRODUCTION
In recent years, building large-scale generalist models capable of solving multiple tasks has become a dominant research focus in natural language processing (NLP) and multi- molliitmddlltyy (Yue et al.) 2024]. Their success is largely driven by the scaling law (Kaplan al. 2020). which posits that increasing model size and data leads to improved performance. However, similar scaling trends have not been extensively observed in reinforcement learning (RL).
Unlike in vision and language domains, RL has traditionally favored smaller models tailored to single tasks or multiple tasks within the same environment. Concerningly, previous studies have shown that scaling model capacity can lead to instabilities or performance degradation (Kumar et al. 2020al Ota et al., [022 Sokar al. 2023). While some efforts have been made to scale offline RL across multiple tasks (Lee et al., 2022 Xu et al., 2022), they predominantly rely on supervised learning (SI.) approaches, such as conditional behavior cloning, rather than temporal difference (TD) learning, and heavily rely on large amounts of expert trajectories. Kumar et al. (2023) scaled offline Q-learning using ResNet-based representation network with separate Q-heads for each game, but this approach only learns generalizable representations and still requires substantial data and gradient steps to adapt to unseen games due to reset of Q-heads during fine-tuning. Therefore,
"qqual corresponding author

---

Published as a conference paper at ICLR 2025
sealing TD-based offline for simultaneous generdd general-purpos representation and decision remains a critical challenge. While Hansen et al. (2024) attempted to address this challenge for continuous control tasks with low-dimensional proprioceptive states using model-based RL, it lacks generalization due to heterogeneous proprioceptors across task domains.
Meanwhile, world model has decoupled from model-based RL and evolved into distinct research area within computer vision and multi- moddlityy primarily focusing on conditional video genera- tion (Hong et al., 2022; Blattmann et al. 2023 Yu et al. 2023] Yang al. 2024] Bruce al. 2024). Notably, SORA (Brooks et al., 2024) has demonstrated superior generalization performance as a world simulator through large-scale training of generative models on time-series image data. This motivates our investigation into a compelling question: "Can image observation-baat based world model scale offline RI. across multiple tasks while enhancing generalization to diverse unseen tasks?"
To address this question, we introduce JOWA: Jointy-Optimized World-Action model, an offline model-based RL agent pretrained across multiple visual games with approximately 6 billion tokens data. Crucially, JOWA unlocks scaling trends and achieves sample-efficient adaption to novel games. By utilizing a shared transformer backbone for both the world model and Q-value criticism, JOWA learns both generalizable representations and decision-making skills. This architecture allows the transformer to absorb gradients back Sack-propagated from both world modeling and TD losses, enabling joint optimization. The world modeling loss acts a regularizer, stabilizing TD-learning for large models. Additionally, we propose a provably efficient and parallelizable planning algorithm to com- pensate for the Q-value estimation error, allowing for consistent identification of the optimal policy at inference time and sample-efficient transfer to novel games.
To evaluate the performance of JOWA, we train single model play 15 Atari games, similar to Lee et al. (2022) but using a reduced yet sufficient dataset of 10M transitions per game-termed as the low-data regime to highlight the data efficiency. This setup presents a significant challenge due to the unique dynamics, visuals, and agent embodiments of games. To further test JOWA's generalization, we perform offline fine-tuning on 5 unseen games, using minimal fine-tuning data.
Our contributions are threefold: First, we introduce JOWA, an offline model-based RL method capa- ble of training a single high- performinf generalist agent across multiple Atari games. JOWA attains 78.9% level performance on pretrained games using only 10% of the original dataset (Agar- wal et al. 2020), outperforming existing state-of-the-art large-scale offline RL baselines by 71.4% on averange. Second, we demonstrate that JOWA unlocks scaling trends, with performance improv- ing as model capacity increases. Third, JOWA enables sample-efficient transfer to diverse unseen games with 64.7% DON-normalized score using only 5k transitions per game, surpassing baselines by 69.9% on average. Our ablation studies highlight the significance of two key design features of JOWA: joint optimization and planning, along with other training choices. We release all the training, evaluation, and fine-tuning code, as well as model weights, to support future research.
2 RELATED WORK
Offline Reinforcement Learning. Offline RL algorithms learn policy entirely from the static offline dataset without online interactions. Model-free offline RL incorporates conservatism to miti- gate extrapolation error (Jin et al.] (2021) primarily through policy constraints (Fujimoto et al. 2019)] Kumar et al., 2019: Fujimoto & Gu 2022 Cheng et al.] 2024 Liu et al., 2024 2025) and value regularization (Kumar et al.. 2020b). Model-bbaed offfine RI. approximates the environment using world models and performs conservative policy optimization (Yu et al. 2020 2021). While these works focus on single-task settings, our work explores scaling offline model- based RL across di-
verse, challenging multi-task Atari games (Lee et al. 2022 Kumar et al. 2023 Wu et al. 2024)
aiming for sample-efficies transfer to novel games.
Mutli-Task Reinforcement Learning. Multi-task reinforcement learning (MTRL) aims to learn a shared policy for diverse tasks. A common approach is to formulate the multi-task model as lask-condition, such as language-conditinn conditional RL (Plappert et al.) 2018). In multi-task offline RL, conditional sequence modeling approaches based on decision transformer (Chen et al.] 2021) or diffusion model (Janner et al., 2022) typically rely on large amounts of expert trajectories (Xu et al.] 2022 Lee et al. 2022] Wu et al. 2024; Hu et al., 2024: He et al.) 2024). Beyond that, FICC (Ye et al.] 2022) pretrains multi- task representation and dynamic models with action-free videos and then fine- tene-tuns& model-based
tasks (Ahn et al.) 2022 Jang et al. 2022) and goal-
2

---

Published as a conference paper at ICLR 2025
agent each task for fast adaptation. Scaled-QL (Kumar al. 2023 scales offline Q-learning using a shared feature network across tasks with separate Q-value heads for each task. Our work advances offline TD learning to multi-task settings without task-specific Q-value heads through a jointy-optimized world-action model. Table I compares experimental environments and open- source status of multi-task offline RI. algorithms. Following [Lee et al. (2022); Kumar et al. (2023); Wu et al. (2024), we focus on the multi-game regime, which presents greater challenges due to high-dimensionall observations and diverse, stochastic environment dynamics.
Table 1: Comparison of methods in multi-task offline RL. and represent two training paradigms of agents, conditional BC and TD-learning, respectively.
Experimental enviruament
Open source
Method
Observation
Action
Dynumics
Check-
Benchmark
Train
Eval
space
task
across tasks
points
TD-MPC2 Hansen ct al. 2024
HarmoDT (Hu et al. 2024
or DMControl
Meta- World
state
continous
deterministic
same
similar
/
d
K
FICCA Ye et al. 2022
x
x
x
MGDT (Loc et a.. 2022
x
/
d
Elastic DT Wi al. 2024
Atari 2600
pixels
discrete
stochastic
diverse
x
Scaled-QL Kumar ct al. 2023
X
JOWA (ours)
PRELIMINARIES AND PROBLEM SETUP
3.1 ONLINE DISTRIBUTIONAL RL (C51)
In distributional RI., the distribution over returns replaces the Q-value the Bellman optimality equation. The Q-value is the mean of the value distribution Z that can be computed through distributional Bellman optimality operator (Bellemare et al. 2017).
T 7(s _ R(s + + 7Z8, arg max Q(''aa)))
where the formula Y _ U denotes equality of probability laws, that the random variable Y is distributed according to the same law U. The C51 algorithm (Bellemare et al.) 2017) models 7 (s, a) using discrete distribution supported on N fixed atoms 2 S S nN uniformly spaced over predetermined interval. Given current value distribution, C51 applies a projection step to map the target distribution onto its finite element support and optimizes as follows:
LID = DKL Zo (s, a) (o) a))
(2)
3.2 VALUE REGULARIZATION BASED OFFLINE RL (CQL)
To be conservatism on unseen actions, CQL (Kumar et al. 2020b) introduces a regularizer to the TD-loss, which minimizes Q-values for unseen actions while maximizing Q-values for actions in the dataset to counteract excessive underestimation. The loss function for CQL is given by:
log Eexp [exp((((ssd))))
Leqr. EEND
Ex [Qa(s, + LTD
(3)
3.3 PROBLEM SETUP
We consider a multi-task offline RL problem: given a static dataset of transitions D {(st, de, Te, de 1+ collected from various environments with arbitrary behaviour polices, our goal is to learn single policy that maximizes the expected return R2 = Σx>0 A on all con- sidered environments and can be efficiently fine-tuned to new tasks, where 7 is the discount factor.
Considering the computational budget, of-distribution (OOD) experiments. The whole training process took around 12 days on A A100 GPUs.
15 Atari games for pretraining and 5 games for out-
FICC only open-sourced the world model pretraining code without the agent fine-tuning code. "Salled-QI. released the preliminary code which not run-able out-of-the-box.
3

---

Published as a conference paper at ICLR 2025
2
Q1
22
2
Po
Po
Rg
Po
Pr
Pd
Po
Po
hq
P+
Pd
Transformer Backbone
Legend:
a2
a2
Network
Transition
enc
dec
enc
dec
Inner output
02
0₂
62
Final
Figure I: Architecture of JOWA. We use a shared transformer backbone for both world modeling and Q-value criticism to enable joint optimization. VQ-VAE tokenizes images into visual tokens. The sum of vocabulary embeddings, position embeddings and task embeddings forms the input embeddings space for the transformer backbone.
Our offline dataset is derived from the the DEN-Replay dataset (Agarwal et al. 2020). which con- sists of 84 84 grayscale images observations and a full action space with 18 discrete actions.
JOINTLY-OPTIMIBE WORLD-ACTION MODEL
In this section, we first detail the architecture of JOWA and the loss functions for joint optimization in section[4.] Next, we introduce the provably efficient and parallelizable planning algorithm em- ployed to compensate for the Q-value estimation error in section 4.2 Finally, we present the overall training pipeline in section
4.1
WORLD-ACTION MODEL
4.1.1 ARCHITECTURE
Figure [ illustrates JOWA's architecture, which uses a transformer backbone (Vaswani 2017) to simultaneously learn world dynamics and Q-values across environments. This dual capability is achieved through distinct prediction heads that process the transformer's ' output embedding e. The world dynamics modeled via supervised learning using three heads: next observation token predictor Po, reward predictor Pr. and termination predictor Pa- The Q-values head lig learns the Q-function through TD-learning, based on implicit representations of historical trajectories. In the following sections, we refer to the (transformer, Po. pa} components as the "world-part" with parameters and the (transformer, h2) components as the "action-part" with parameters Ø.
We use VQ-VAE, a discrete autoencoder, to tokenize image observations, representing high- dimensional images as sequence of K tokens. The VQ-VAE trained with an equally weighted combination of L, reconstruction loss, commitment loss (Van Den Oord et al. (2017), and perceptual loss (Esser et al., 202I). Then the transformer operates on the interleaved observation and action tokens represented as (=),... ... = ag..... 2/,,... .., K aL.) where I. is the maximum timesteps of the trajectory segments. For multi-task learning, we incorporate learnable task embeddings for both observation and action tokens.
4.1.2 TRAINING OF WORLD-PART MODULE
At each timestep I, the world-part module models the following distributions:
Dynamics predictor:
* ~ po f(a5-11, 2 (ase-1,u))
(4)
Reward predictor:
R, ~ pr (Fr [f
(5)
Termination predictor: d, palde f(zg()
(6)
4

---

Published as a conference paper at ICLR 2025
where f represents the transformer backbone, is the task ID index the corresponding task embedding, h E {1,... ,K), and E {1,-.., ,L}.
To unify the category of loss functions for balanced training (Vandenhende al. 202 convert scalar rewards to ternary quantities {-1, 0, 1} using the sign function. This allows all three predic- tors be optimized classification problems by cross-entropy loss. Given Z-timesteps segments sampled from the offline dataset, the loss function for the world-part module is formulated as:
1
woorld (0)
Inpo f as)-1,
Inp, (re ((zs) Inpalde f(zseass,u))
4.1.3 TRAINING OF ACTION-PART MODULE
We use CQL for offline TD-leamnin during pretraining. To ensure training stability and enhance scaling performance (Farebrother et al. 2024), we employ distributional TD-error (Bellemare et al.) 2017) instead of the mean-square TI)-error maintaining consistency with the L-worid loss category.
The action-part module computes the return distribution for all actions given an observation and historical information. The value function (ot, ae) is the mean of Z(or, 0z). Then the loss function for the action-part module is formulated as:
(8) where T is the distributional Bellman optimality operator defined in Equation and 2- the target distribution computed through a target Q-values head li--. We a 0.1 experiments.
Lattion (x) [log Ca exp (Q(o,, a))) Q(or, at) + DEL (T'Z (()), az))
Therefore, the joint optimization objective for the world-action model is formulated as:
L(0, Φ) =BL + (action())
with the coefficient / We set 3 0.1 in experiments.
Both Lwarld and Laction back pack-prop gradients to the transformer. Previous works observed that TD methods suffer from greater instability with larger model size (Kumar et al.) [020a Sokar al., 2023). However, through jointly optimizing the world-action model, L serves as a regularizer to stabilize TD-learning large-scale models. Moreover, the world-part module enables planning decision time for optimal inference and sample-efficien transfer, detailed in the following sections.
4.2 PARALLELIZABLE PLANNING INFERENCE TIME
The world-part module enables planning at decision time compensate for inaccurate Q-value estimates, allowing JOWA to consistently search out the optimal policy. We model this process as tree search problem and present a practical, parallelizable search algorithm.
Given an estimated optimal Q-value Q*, a learned world model with dynamic predictor P and reward predictor f, our objective is to find the optimal action aj maximizing the ground-truth optimal Q-value Q", starting from state 8g. To do so, we rewrite the Bellman optimality equation as:
Qs ao) max To"
E
(r((s, n max Q" (sn, (s)r,a)) Σ 1=
(10)
ag
where TQ- is the policy induced by the optimal Q-function. For Equation 10 the optimal policy is the greedy policy based on Q". The proof is provided the Appendix
To derive the search objective function, follow three steps: (i) replace the ground-truth functions in the right side of Equation (1O) with estimated or learned functions. (ii) leverage the estimated
optimal Q-function Q" to reduce the policy space for search, restricting the actions to those with top-K highest Q-values. Denote the constrained policy space as H.... where YI E II . Vs
S, Va & (op-K(("(ss (Q((sss))))) (s.+)) )), we have (als) (iii) maximize ag both sides of Equation

---

Published as a conference paper at ICLR 2025
10 to find the optimal initial action, considering the restriction in the second step. Finally, the resulting objective function formulated as:
max xEllo
E P
Σ (()))
+7 max Q" (sH, an) an
(11)
$1
1
income-to-go
income-to-date
Detailed derivation is provided in the Appendix B Then we show the error bound of search-based optimal Q-function in Theorem 4.1 with the formal theorem and its proof in Appendix
Theorem (Informal) 4.1. Denote the search-based optimal Q-function as [(s, a). Assume the learned reward function to be Lipiichc and the estimated optimal Q-function Q to be L-q Lipschitz. Assume the estimation errors of the learned state transition, reward, and Q-value are bounded by ea. Er, EQ respectively. Then we have the error between search- baace optimal Q-value f(s a) and ground-truth optimal Q-value Q" (s, a) bounded as:
1- 7 - er + (2-2- =ca) e, + +7
(1)(,, - Q=
(12)
Under the following condition, the search-based optimal Q-function has upper error bound no greater than the estimated optimal Q-function Q":
Gr ( ? L, n LQ 1-7
1
L, LQ 7 1-9 1 1-1
Es SEQ
(13)
1-7 1 7
interact with the imagined MDP induced by the learned world model using the constrained policy llg- for H steps, starting from 8p. We then compute the total income (income-to-datt plus income-to-go) for all leaf nodes to identify the optimal path. The first edge of this path is the optimal initial action.
To optimize objective function [1,,
We implement this search using beam search, an efficient decoding algorithm common NLP. At each timestep, retain only K states with the top-K total income values. The horizon // and beam width K are syper-parametery The algorithm calls the world model ( 1) + K times but only takes H times as long as the forward propagation of the world model due parallelizability across K beams.
with K
H degenerating to a Q greedy policy.
4.3 TRAINING PIPELINE
Our multi-task offline RL pretraining consists two stages:
Stage Sample trajectory segments from datasets. Train the VQ-VAE tokenizer using image observations. Train the world-part module using segments with loss 7 for M1 steps.
Stage 2: Freeze the VQ-VAE tokenizer. Sample segments and jointly optimize the world- action model with loss function ( for M2 steps.
We employ this two-stage training approach to stabilize and accelerate the overall training process. our pretraining experiments, we set M, = 250k and M2 = 1.5M, totaling 1.75M gradient steps.
EXPERIMENTS
We design experiments to answer the following questions: (1) How does JOWA perform multi-games in low-data regime? (2) Can JOWA effectively leverage higher model capacity? (3) Does the pre-trained JOWA sample-efficiently transfer games?

---

Published as a conference paper at ICLR 2025
5.1 EXPERIMENTAL SETUP
Dataset. We use the Atari dataset from Agarwal et al. (202)), which contains 50M transitions from each of separate training runs. Due to the prohibitive long training time on full games, we select subset of 20 games, maintaining the difficulty distribution of full games defined by Gulcehre et al. (220)) These 20 games are introduced in Appendix [D] 15 of those games are used for training, and games are held out for OOD generalization experiments. Following Lee et al. (2022), we use data from out of 5 training runs, To investigate performance in low-data regime, we uniformly draw 10% of transitions at random, as per per Aggaraa et al. (2020), resulting in 10M transitions per game.
Training and Fine-tuning. We implement our world-action model based on GPT-2 (Brown et al, 2020). We train three JOWA variants: JOWA-150M (150M parameters), JOWA-70M, and JOWA- 40M. We set the number of visual tokens K to 36, resulting in a total of approximately 6B tokens in our dataset. The sequence length L is set to 8. We pretrain all JOWA models A A100 GPUs for 1.75M steps. For fine-tuning, we train models for 50k gradient steps with 5k transitions.
Evaluation and Metrics. During evaluation, we enable planning for JOWA, setting the planning horizon H to 2 for all games and adjust the heam width K based on the valid action space size of each game. See Appendix E.3 for the detailed evaluation protocol. We measure performance using human normalized scores (HNS) (Mnih et al. 2015). i.c. (score score (scorehuman scCOeeandom) To create an aggregate comparison metric across all games, we use inter-guartile mean (1QM) of numan-normalized scores, following evaluation best practices proposed in (Agarwal et al. 2021).
More details on hyperparameters, network architecture, algorithm implementation, and protocols for fine-tuning, and evaluation are provided in the Appendix F
5.2 BASELINE METHODS
We compare JOWA with the following baselines: (i) Multi-Task BC (MTBC), Our method can nat- urally be reduced to transformer-based multi-task behavioral cloning agent. We train MTBC mod-
els at 3 scales: 34M, 65M, 120M parameters. (ii) Multi-Game DT (MGDT). et al. (2022) used
{oe+i, Retis dz+i, Te++ return distribution to induce expert-level actions during evaluation. We train MGDT models at scales: 40M and 200M parameters. (iii) Elastic DT (EDT). Based MGDT, Wu et al. (2024) removed the reward in the input sequences and dynamically selected history length to address chal- lenges of trajectory stitching. We train EDT based the architecture configuration of MGDT- 200M. (iv) Scaled-QL. (SQL). Kumar et al.] (2023) scaled offline Q-learning through pre-trained and shared representation network and separate Q-values head for each game. We reproduce the largest Scaled-QL with 80M parameters using 101 network. (v) FICC. Ye et al. (2022) pretrained the dynamic model with videos and then online fine-tuned the agent on each game. To obtain an offline multi-task policy, we pretrain FICC for 0.5M steps and then primarily fine-tune the Q-function on 15 games all at once, the same as JOWA's second pretraining stage, for 1.25M steps. We replace all residual blocks in FICC's official code with ResNet- blocks, resulting in FICC-85M. For fair comparison, methods use the same batch size of 512 and 1.75M gradient steps. The implementation details of all baselines are provided in the Appendix E.
input sequences, where it is the target return, and expert-conditioned
5.3 How DOES JOWA PERFORM ON MULTI-GAMES IN LOW-DATA REGIME?
We summarize our main results in Table 2 This table shows the performance of JOWA alongside all best performing sizes of baselines trained with 10% subsampled dataset. MTBC, MGDT, and EDT represent (conditional) behavior cloning methods, while Scaled-QL. FICC, and JOWA repre- sent Q-learning methods. Despite the constraints on the amount of data, JOWA achieves superior performance with comparable or fewer parameters than baselines. Specifically, JOWA-150M attains an IQM human-normalized score of 78.9% in low-data regime, surpassing MGDT-200M (49.8%) and EDT-200M (50.2%), despite these two models having more parameters and utilizing data aug- mentation during pretraining. JOWA-70M achieves an IQM human- nomaa--nrmmalieed score of 47.6%, outperforming Scaled- -L-MMMMQLLOMM (42.0%) and FICC-85M (43.3%). JOWA maintains its performance advantage when comparing median human-normalizz normalized scores. These results demonstrate JOWA's sample efficiency in learning from heterogeneous offline data.
7

---

Published as a conference paper at ICLR 2025
Table Returns on the 15 Atari games for best performing sizes of multi-game models trained with the 10% subsampled dataset. Bold numbers indicate the top methods.
MTBC 120M
MGDT
EDT
SQL 80M
FICC
JOWA
Game
Random
Human
200M
200M
85M
40M
70M
150M
Assault
222.4
742.0
1203.5
1741.5
1915.5
2292
925,9
1578.3
1733.9
2302
Atlantis
12850
29028.1
37812.5
2565750
2108166.7
49100
86250
41662.5
570862.5
2690387.5
BeamRider
363.9
16926.5
786
6011.3
4149.8
7023,5
6822
880.3
2547.4
3498
Berzerk
123.7
2630.4
655
444.5
350
312.5
400
395
441.9
739
Carnival
0
3800
5560
2610
3678.8
1940
2820
4685
40770
5316
Centipede
2090.9
12017
4046.6
4604
3389.9
3650.3
3742.2
5669.3
44755.6
4677
ChopperConmannt
811.0
7387.8
256.3
3300.8
3443.8
853.8
2835.6
3118.8
2568.8
3812.5
DemonAAlack
152.1
1971
2611.3
6549.4
3455.7
7936.5
5806.4
1233.8
4584.4
3547.8
Name ThisGimee
2292.3
8049
6045
6610.5
7060
6461.7
6236
7335.6
12706.9
11421
Phoenix
761 4
7242.6
1446.9
5120.5
5320
4961.7
3814.5
1608.8
5065
5348
Seaquest
68.4
42054.7
120
2720
3160.4
650
1760
1033,1
1490
2725
Spacelmadeess
148
1668.7
605
742.5
$13.8
714,4
641.2
694.7
969.1
744.7
SturGunner
664.0
10250
1493.8
8625.5
9550
4728.6
4936.4
5737.5
21231.3
18150
TimePilot
3568
5229.2
762.5
3866.7
2812.5
4072.7
4166.7
3268.8 8
3831-3
3669
Zaxxom
32.5
9173.3
462.5
325.5
312.5
0
225
2163
#Superhuman
0
N/A
4
3
3
6
Median HNS
0.400
0.360
0.456
IQM HNS
0.000
1.000
0,329
0.498
0.502
0.420
0.438
0.371
0.476
0.789
Next, we present experiments highlighting JOWA's scaling and generalization capabilities. These two key properties derived from multi-game pretraining align with recent advancements in SL.
5.4 HOW DOES JOWA SCALES WITH MODEL SIZE?
Scaling law depicts the positive correlation between model capacity and performance. Following Lee et al.
2002 Kumar al. (2023), we investigate JOWA's abil-
Human-Norrallize IQM
ity to leverage higher capacity architectures. For com- parison of scaling trends, we establish baselines: (i) MTBC scaled to 34M, 65M and 120M parameters. (ii) MGDT with 40M and 200M parameters. We scale JOWA by increasing the size of the transformer backbone and Q-values heads, resulting in 3 variants with 40M, 70M, and 150M parameters. Figure shows the scaling trends for the algorithms, which demonstrates that JOWA's performance reliably increases the model size grows. Although previous works observed that TD-learning suf- fer from greater instability with larger model size (Lee) iet al. 2022 Sokar et al. (2023). JOWA scales TD-learning through a scalable transformer architecture and an auxil- iary regularization loss Lwarlu to stabilize the TD-learning in large models. Notably, JOWA exhibits the steepest scaling curve among all algorithms, highlighting the great scalability potential of offline model-based RL.
MTBC
60%
50%
60
230
Parameters (x1 Millaan)
Figure 2: Scaling trends for different al- gorithms on the training set games.
Table 3: Fine-tuning performance on unseen games using 5k transitions, measured in terms of DQN- normalized score, following Lee et al. (2022). See Table in Appendix for raw scores.
MTBC
MGDT
EDT
SQL
FICC
JOWA
JOWA-1500 (scratch)
DNS
120M
200M
200M
80M
85M
40M
70M
150M
Mean
0.164
0.422
0.430
0.360
0.543
0.504
0.576
0.647
0.196
Median
0.215
0.354
0.325
0.284
0.565
0.512
0.715
0.615
0.173
IQM
0.205
0.377
0.380
0.355
0.575
0.498
0.603
0.647
0.181
5.5
CAN JOWA SAMPLE-BIFICIENTLY TRANSFER TO NEW GAMES?
Rapid adaptation to downstream tasks is a natural and well-motivated benefit of pretraining. To investigate this question, we fine-tune pretrained agents on held-out games using uniformly sub- sampled 5k expert-level transitions (from last 20% of DQN- DON-Replay (Agarwal et al. (0200) per game as the benchmark. These tiny amounts of transitions corresponding to approximately trajectories from expert-level DON-Replay per fine-tuned game average, which is similar to the settings of
8

---

Published as a conference paper at ICLR 2025
few-shot learning and is extremely challenging. For comparison, fine-tune the largest agents of baselines alongside a no- po-pretrained baseline, JOWA-150M (scratch) trained on each held-out game from scratch. All pretrained agents and JOWA- (scratch) are fine-tuned for 50k and 500k gra- dient steps, respectively. Detailed fine-tuning protocol is shown in Appendix E.2 We report results
in terms of DON-normalized score (DNS), following [ee al. (2022). in Table
The results show that the fine-tuned LOWA-150M. attains 64.7% IQM DNS across held-out games, outperforming baselines by 69.9% on average. Moreover, comparing the performance of 3 JOWA variants, the scaling trend on IQM DNS still holds for novel games after fine-tuning. These results underscore JOWA's capacity for rapid and sample-efficiett transfer to novel games, highlighting the efficacy of its learned general-purpose representation and decision-make makiing capabilities.
5.6 ABLATION STUDY
In this section, we conduct a series of ablation studies to evaluate the impact of key design choices in JOWA, including planning, different training losses, and usage of synthetic data. More experiments can be found in Appendix F These experiments aim to provide empirical evidence to support our designs and offer valuable insights for future research this domain.
Table 4: Performance on the 15 Atari games for variants of JOWA evaluated with or without planning reported in terms of the IQM HNS.
Table 5: Results on 7 games for JOWA- 150M evaluated with different planning algorithms. See Table 19 raw scores,
JOWA-1500
JOWA-70M
JOWA-40M
w/o planning
MCTS
Ours
without planning with planning improvement (1)
50.8%
45.1%
27.1%
FPS Mean HNS
10.8
0.12
1.26
78.9%
47.6%
37.1%
-2%
22, 156
37.8%
+28.1%
+2.5%
+10.0%
IQM HNS
13.4%
23.7%
Effeets of planning at decision time. We evaluate our 3 variants models with or without plan- ning and report the IQM HNS across 15 pretrained games in Table 4 Observe that the addition of planning improves 13.5% performance on average and 28.1% on maximum for JOWA-150M. Even without planning, JOWA still exhibits scaling trends with increasing model size, but a slower rate.
To further demonstrate the efficiency of our planning algorithm, we compare it with MCTS. Be- cause Muzero-style (Schrittwieser et al.) 2020) MCTS requires access to the V-function and pol- icy networks π while JOWA only estimates the optimal Q-function, we use max, Q(s, a) as the V-value (ss) and use an energy-based policy to compute the action probability, i.e., (-]ss) softmax(Q( 1/t), where t is the temperature. We conduct a grid search on the implementation choices and hyperparameters for MCTS and show the details in Appendix [ We report the frames per second (FPS) and HNS on 7 games, where bare JOWA- 150M performs poorly. in Table 5 Ob- serve that ours not only 10x faster than MCTS but also exceeds MCTS by 71.0% on Mean HNS. Moreover, we empirically observe that MCTS is highly sensitive to the temperature and max depth of the tree while the long exccution time makes it inconvenient to tune hyperparameters.
The subsequent ablation experiments require training from scratch in multi-game regime. For time- saving, we consider a subset of games in the following experiments: Assault, Carnival, Centipede, NameThis whw.hhisaaeeee sGame, Phoenix, Space Invaders. We train all models with 10% randomly subsampled data per game for 1M gradient steps and fix the parameter size to 40M.
Table 6:
The mean, median, and IQM human-normallan normalized score on the 6 Atari games for various
training choices. See Tabel 20 Appendix
F for raw scores.
Different training losses sg (g([etiio))
Synthetic data in pretraining
Game
Origin
No CQL
No Lwond
MSE
Mcan HNS
1.183
0.613
0.917
0.293
0.189
0.448
Median HNS
1.078
0.696
0.489
0.304
0.118
0.452
IQM HNS
1.123
0.637
0.659
0.307
0.126
0.464
Different training losses. We train models with each of the following 5 loss functions to investi- gate the impact of each loss term: (i) original loss defined in Equation (9). (ii) No regularizatio in Lacuion which means no conservative constraints, (iii) No Lworls, which means training model- free transformer- based critic network using Caction. (iv) gradients of Laction not back- pooc-propgggted to transformer backbone, denoted La(Caition) for short, which means optimizing the world model

---

Published as a conference paper at ICLR 2025
and critic network separately, (v) mean square enror (MSE) instead of distributional TD-loss. The results of different training losses are shown in the third main column of Table 6
Observe that the sg(Cactioo) and MSE loss fail to train qualified multi-game game agents. We empirically observe that agents trained with MSE TD-loss always over-optimize the CQL regularization term regardless of the value of coefficient a (tested with E {0.01, 0.05, 0.1 }), resulting in extremely high Q-values for in-domain state-action and low Q-values for OOD actions. The fail of sg (faction) underscores the importance of foint-optimization in the world-action model. Comparing the original loss with the no Lworl configuration demonstrates that scaling model-based multi-game RL is more efficient than model-free RL. Additionally, the CQL conservatism regularizer is also necessary in multi-game pretraining. Overall, the original loss outperforms all variant losses by large margin, indicating the effectiveness of every loss terms.
Effects of synthetic data in pretraining. By default, we do not employ planning synthetic data during pretraining due to the significant increase in training time (approximately 10x slower, extending training months). For this ablation study, we enable synthetic data in pretraining for JOWA-40M and train on 6 games for I M steps, which takes 31 days. Specifically, we sample batch of 4-step segments and interact in the imagined MDP with e-greedy policy to synthesis the last steps. Then we optimize JOWA with half real data and half synthetic data using COMBO loss as Laction- The results are shown in the last main column of Table
Suprisingly, we observe negative gains from the usage of synthetic data. We hypothesize twofold reasons: (i) accumulation of inference errors over steps causes later synthetic steps to deviate signif- icantly from the ground-truth distribution, (ii) COMBO's pver-eenalization of Q-values for unseen state-action pairs results in overly conservative agents. Due to the unbearable training time, further investigation into more effective methods of utilizing synthetic data is left for future work.
6 CONCLUSION
In this work, we introduce JOWA: Opttil-d World-Action model, a single offline model- based RL agent capable of playing multiple Atari games. JOWA uses a shared transformer backbone for both the world modeling and the Q-value criticism, enabling joint optimization. We propose a provably efficient and parallelizable planning algorithm to consistently identify the optimal policy during inference. As we hoped, by increasing model parameters, JOWA unlocks scaling trends in performance and exceed prior large-scale offline RI. methods in multi-games regime. Furthermore, by training a large-capacity model on diverse set of games, we show that JOWA can sample- efficiently adapt to novel games, leveraging its generalizable world model for planning. Our ablation studies validate the efficacy of joint optimization and our planning method, while also demonstrating that scaling offline model-based RL is more efficient than offline model-free RL. To facilitate future research, we will release all training and evaluation codes, along with pretrained model checkpoints.
Limitations. Due to the computation resources required, we could not experiment on full Atari 2600 games with complete datasets (Agarwal et al. 2020 as this would take 2 months for training. Additionally, while using synthetic data to train agents is common in single-task online model- based RL, we observe it yields negative gains in multi-game RL pretraining. Therefore, how to effectively use synthetic data for multi-task pretraining is an interesting direction for future work. Lastly, al- though we observe the scaling trend on IQM aggregated scores, the raw scores in some games against this trend, i.e., increasigg model parameters leads to decreasing performance in specific games. We hypothesis that this inconsistence is an inherent problem of offline RL. Unfortunately. to knowledge, neither increasing model size nor extending training iterations can solve this issue.
ACKNOWLEDGMENTS
We are grateful to the anonymous reviewers and editors for their insightful suggestions. This work was partially supported by the National Science and Technology Major Project (2022ZD011710) the National Natural Science Foundation of China under Grants 62271485 and 62303462, Bei- jing Natural Science Foundation under grant L241016 and L233005, Chongqing Transportation Sichuan Key-Area Research and Development Program 2024YFHZ0011, the Provincial Key Research and Development Program of Zhejiang (project num- ber: 2022C01 129), Ningbo International Science and Technology Cooperation Project (2023H020).
Technology Project (CQJT-CZKJ2022
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Rishabh Agarwal, Dale Schuurmans, and Mohammad Norouzi. An optimistic perspective on offline reinforcement learning. In International Conference on Machine Learning, pp. 104-114, PMLR, 2020.
Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C Courville, and Mare Bellemare. Deep reinforcement learning the edge of the statistical precipice. Advances in Neural Informa- tion Processing Systems, 34:29304-299 2922000-99320 2021.
Michael Ahn, Anthony Brohan, Noah Brown, Yevgen Chebotar, Omar Cortes, Byron David, Chelsea Finn, Chuyuan Fu, Keerthana Gopalakrishnan, Karol Hausman, et al. Do as i can, not as say: Grounding language in robotic affordances. arXiv preprint arXiv:2204.01699 2022.
Mare G Bellemare, Will Dabney, and Rémi Munos. A distributional perspective on reinforcement learning. International Conference on Machine Learning, pp. 458. PMLR, 2017.
Andreas Blattmann, Tim Dockhorn, Sumith Kulal, Daniel Mendelevitch, Maciej Kilian, Dominik Lorenz, Yam Levi, Zion English, Vikram Voleti, Adam Letts, et al. Stable video diffusion: Scaling latent video diffusion models to large datasets. arXiv preprint arXiv:2311.1512 2023.
Tim Brooks, Bill Peebles, Connor Holmes, Will DePue, Yufei Guo, Li Jing, David Schnurr, Joe Taylor, Troy Luhman, Eric Luhman, Clarence Ng. Ricky Wang, and Aditya Ramesh. Video generation models as world simulators. 2024. URL httpp:: /ttttp:::/oppnaii.. com/ research/eec/eeeeee
Sideo-generatiin modeea
world simulators
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 1877-1901. Curran Associates, Inc., 2020.
Jake Bruce, Michael D Dennis, Ashley Edwards, Jack Parker-Holder, Yuge Shi, Edward Hughes, Matthew Lai, Aditi Mavalankar, Richie Steigerwald, Chris Apps, et al. Genie: Generative inter- active environments. In International Conference on Machine Learning. 2024.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Picter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in Neural Information Processing Systems,
2021.
Jie Cheng, Gang Xiong, Xingyuan Dai, Qinghai Miao, Yisheng L.,, and Fei-Yue Wang- Rime: Robust preference-hased reinforcement learning with noisy preferences. arXiv preprint arX1v:2402.17257, 2024.
Tri Dao. Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with io-awareness. Advances Neural Information Processing Systems, 15:16344-16359, 2022.
Patrick Esser, Robin Rombach, and Bjom Ommer. Taming transformers for high-resolution image synthesis. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recog- nition, pp. 12873-12884 2021.
Jesse Farebrother, Jordi Orbay, Quan Vuong, Adrien Ali Taiga, Yevgen Chebotar, Ted Xiao, Alex Irpan, Sergey Levine, Pablo Samuel Castro, Aleksandra Faust, et al. Stop regressing: Training value functions via classification for scalable deep rl. arXiv preprint arXiv:4403.03950 2024.
Scott Fujimoto and Shixiang Shane Gu. A minimalist approach offline reinforcement learning. Advances in Neural Information Processing Systems, 34:20132-20145, 2021.
Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In International Conference on Machine Learning, pp. 0052-2062. PMLR, 2019.

---

Published as a conference paper at ICLR 2025
Caglar Gulcehre, Ziyu Wang, Alexander Novikov, Thomas Paine, Sergio Gómez, Konrad Zolna, Rishabh Agarwal, Josh S Merel, Daniel Mankowitz, Cosmin Paduraru, et al. RI unplugged: A suite of benchmarks for offline reinforcement learning. Advances in Neural Information Process- ing Systems, 33:7248-7259, 2020.
Nicklas Hansen, Hao Su, and Xiaolong Wang. Td-mpe2: Scalable, robust world models for contin- uous control. In International Conference on Learning Representations, 2024.
Haoran He, Chenjia Bai, Kang Xu, Zhuoran Yang, Weinan Zhang, Dong Wang, Bin Zhao, and Xue- long Li. Diffusion model an effective planner and data synthesizer for multi-task reinforcement learning. Advances in Neural Information Processing Systems, 36, 2024.
Wenyi Hong, Ming Ding, Wendi Zheng, Xinghan Liu, and Jie Tang. Cogvideo: Large-scale pre- video
training for text-to- text-to-video generation via transformers. In International Conference on Learning Representations, 2022.
Shengchao Hu, Ziqing Fan, Li Shen, Ya Zhang, Yanfeng Wang, and Dacheng Tao. Harmodt: Har- mony multi-task decision transformer for offline reinforcement learning. In International Confer- on Machine Learning, 2024.
Eric Jang, Alex Irpan, Mohi Khansari, Daniel Kappler, Frederik Ebert, Corey Lynch, Sergey Levine, and Chelsea Finn. Be-z: Zero-shot task generalization with robotic imitation learning. In Confer- ence on Robot Learning, pp. 991-1002. PMLR, 2022.
Michael Janner, Yilun Du, Joshua Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. In International Conference on Machine Learning, pp. 9902-9915. PMLR, 2022.
Ying Jin, Zhuoran Yang, and Zhaoran Wang. Is pessimism provably efficient for offline rI? In International Conference Machine Learning, pp. 5-84-0044000 5096. PMLR, 2021.
Justin Johnson, Alexandre Alahi, and Li Fei-Fei. Perceptual losses for real-time style transfer and super-resolution. In Computer Vision-ECCV 2016: 14th European Conference, Amsterdam, The Netherlands, October 11-14, 2016, Proceedings, Part 14, pp. 694-711. Springer, 2016.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXh::2001.0836/ 2020.
Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing off-policy q- learning via bootstrapping reduction. Advances in Neural Information Processing Systems, 32, 2019.
Aviral Kumar, Rishabh Agarwal, Dibya Ghosh, and Sergey Levine. Implicit paraeerrrramttiizzaiion inhibits data-efficient deep reinforcement learning. In International Conference on Learning Rep- resentations, 2020a.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems, 33:1 1191, 2020b.
Aviral Kumar, Rishabh Agarwal, Xinyang Geng, George Tucker, and Sergey Levine Offline q- learning on diverse multi-task data both scales and generalizes. International Conference on Learning Representations, 2023.
Shang-HHer Lee, Ofir Nachum, Mengjiao Sherry Yang, Lisa Lee, Daniel Freeman, Sergio Guadar- rama, Ian Fischer, Winnie Xu, Eric Jang, Henryk Michalewski, et al. Multi game decision trans- formers. Advances in Neural Information Processing Systems, 35:27921-27936 2022.
Tenglong Liu, Yang Li, Yixing Lan, Hao Gao, Wei Pan, and Xin Xu. Adaptive advantage divantage-guided guided policy regularization for offline reinforcement learning. In Forty-firs International Conference Machine Learning, volume 235, pp. 31406 31424. PMLR, 2024.
12

---

Published as a conference paper at ICLR 2025
Tenglong Liu, Jianxiong Li, Yinan Zheng. Haoyi Niu, Yixing Lan, Xin Xu, and Xianyuan Zhan. Skill expansion and composition in parameter space. arXiv preprint arXiv:2502.05932, 2025.
Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are sample-efficient world mod- els. International Conference on Learning Representations, 2022.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Belle- mare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning, Nature, 518(7540):559-5555 -33,5999333--0 2015.
Kei Ota, Devesh K Jha, and Asako Kanezaki. Training larger networks for deep reinforcement learning. arXiv preprint arXiv:2102.07920 2021.
Matthias Plappert, Marcin Andrychowicz, Alex Ray, Bob McGrew, Bowen Baker, Glenn Pow- ell, Jonas Schneider, Josh Tobin, Maciek Chociej, Peter Welinder, et al. Multi-goal reinforce- ment learning: Challenging robotics environments and request for research. arXiv preprint arXiv:1802.09464, 2018.
Rylan Schaeffer, Brando Miranda, and Sanmi Koyejo. Are emergent abilities of large language models mirage? Advances in Neural Information Processing Systems, 36, 2024.
Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, al. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588(7839):604-60) 2020.
Ghada Sokar, Rishabh Agarwal, Pablo Samuel Castro, and Utku Evei. The dormant neuron phe- nomenon in deep reinforcement learning. International Conference on Machine Learning, pp- PMLR, 2023.
Richard Sutton. Reinforcement learning: An introduction. A Bradford Book, 2018.
Aaron Van Den Oord, Oriol Vinyals, et al. Neural discrete representation learning. Advances in Neural Information Processing Systems, 30, 2017.
Simon Vandenhende, Stamatios Georgoulis, Wouter Van Gansbeke, Marc Proesmans, Dengxin Dai, and Lue Van Gool. Multi-task learning for dense prediction tasks: A survey. IEEE Transactions on Pattern Analysis and Machine Intelligence, 44(7)3311-3363, 2021.
A Vaswani. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
Shengjie Wang, Shaohuai Liu, Weirui Ye, Jiacheng You, and Yang Gao. Efficientzero v2: Mastering discrete and continuous control with limited data. arXiv preprint arXiv:2403.00564, 2024.
Yueh-Hua Wu, Xiaolong Wang, and Masashi Hamaya. Elastic decision transformer. Advances in
Neural Information Processing Systems, 36, 2024.
Mengdi Xu, Yikang Shen, Shun Zhang, Yuchen Lu, Ding Zhao, Joshua Tenenbaum, and Chuang Gan. Prompting decision transformer for few-shot policy generalization. International Con- ference Machine Leurning, pp. 24631-24645. PMLR, 2022.
Jiazhi Yang, Shenyuan Gao, Yihang Qiu, Li Chen, Tianyu Li, Bo Dai, Kashyap Chitta, Penghao Wu, Jia Zeng, Ping Luo, et al. Generalized predictive model for autonomous driving. In Proceedings of the JEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 14662-14672, 2024.
Weirui Ye, Yunsheng Zhang, Pieter Abbeel, and Yang Gao. Become proficient player with limited data through watching pure videos. In International Conference on Learning Representations, 2022.
Lijun Yu, Yong Cheng, Kihyuk Sohn, José Lezama, Han Zhang, Huiwen Chang, Alexander G Ming-Hsuan
Hauptmann, Ming- Hsuan Yang. Yuan Hao, Irfan Essa, et al. Magvit: Masked generative video transformer. Proceedings the IEEE/CVF Conference on Computer Vision and Pattern Recog- nition, pp. 10499, 2023.
13

---

Published as a conference paper at ICLR 2025
Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Y Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. Mopo: Model-based offline policy optimization. Advances Neural Information Processing Systems, 11114129-141445 2020.
Tianhe Yu, Aviral Kumar, Rafael Rafailov, Aravind Rajeswaran, Sergey Levine, and Chelsea Finn. Combo: Conservative offline model-based policy optimization. Advances in Neural Information Processing Systems, 34:28954-28967. 2021.
Tongtian Yue, Jie Cheng, Longteng Guo, Xingyuan Dai, Zijia Zhao, Xingjian He, Gang Xiong, Yisheng Lv, and Jing Liu. Sc-tune: Unleashing self-consistent referential comprehension in large vision language models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 13073-13083, 2024.
14

---

Published as a conference paper at ICLR 2025
A
PROOF OF EQUATION (10)
Given the initial state-action pair (sp, ao), the bellman expectation equation
Sutton 2018 written
Q (so, au) 81 (-l50
(r(s0, oo) + +Vv*(ss))
H
Σ
rtr(s, ar) + ²
(14)
According to the definition of optimal Q-value: ("(s0, maxx Q (So, ao), we have:
H 7'r(s,, + +
@so Qn) = max x
E
01;"",0g 08
H-1
max
E
Σ 7'r(st, ar) + +7" max (*HH)
P
01," 88 1
H-1
max
E
Σ 7'r(s,, ar) + + (*H))
$1."
01," 08 1M
H-1
max
E
Σ 7'r(s, ar) + + max (sH ah)
(15)
P
we derive the optimal policy æ is the greedy policy selecting actions with the greatest -*-value follows:
For Equation 15
H
Q" (8o: 0n) = max
Σ
7 ar) + max Q'n
P
1=0
H
max
Σ
Ar(S,, ar) +1"1 11 max E s
r(sH-1,4H-1) + max Q (*H,
t=
H-
r'r(se, a²) + H-1 max Q" Q"(SH-1,QH-1) Σ
max
E
(16)
Therefore, have 0 arg max, ("( Similarly, continuing to dynamic pro- gramming on Equation (6). we finally get: = arg max, Q* a), (i = 2, .... H 1). Thus claim that the optimal policy is induced by the optimal value and use the symbol mq- instead of in Equation (15) to obtain Equation (TO).
B DERIVATION FROM EQUATION (10) TO EQUATION (11)
Given the learned dynamic preditor P. reward predictor and estimated optimal Q-value Q first substitute these three functions for the ground-truth functions in the right side of Equation (10):
max 2
IE
Σ ( az) + max Q (SH ah)
(17)
6g
However, due the estimation error between learned functions and ground-truth functions, + Y max, Q- (5141, a)). Therefore, instead
Q* (st, at) is typically not equal to E B [f( of using dynamic programming to derive that the optimal policy is the greedy policy as in Proof have to search solve formula (1p..
15

---

Published as a conference paper at ICLR 2025
the conclusion derived with ground-truth functions make assumptions to reduce the policy space for search. We assume that the optimal policy for formula (17) maps states to actions with top-K highest Q-values. Denote the constrained policy space Q. where Vr e
However,
i Vs E S, Va & top-K(Q" -)), we have T(als) O. we make the following assumption: the optimal policy of formula (17) is in the constrained policy space, q. E Ô Under this assumption, formula
is equivalent to:
E M
Σ ((), + +7" max (sH, an)
(18)
HElL
an
1008
Then we maximize formula assumption, do E top- (). Thus we have the objective for search as:
18
over do to find the optimal initial action. Considering the above
H
max
E
Σ
Afss,, + + max Q (sH,
(19)
P ag oo 1
t=
C UPPER BOUND OF SEARCH-BASEL Q-VALUE ESTIMATION
Let function f(0.. an) be equal to formula (18) and let प.. the optimal policy of formula (88
We have:
[s0; ao)
E f a1 201 minti
Σ t=
7 ('8(s, at) +7 max Q' (SH, an)
(20)
an
We make the following assumption similar EfficientZero-v2 (Wang al. 2024):
Assumption C.1. Assume the state transition, reward, and Q-value estimations error are upper bounded by Eas Er: EQ respectively. The error bound of each estimation is formulated as:
max (E(
E El| s, S c
(21)
max ([(
E [[]]((s) r(s))ll se
(22)
max [E[N]]E[[H[]]
E
-1))(s))))))) (s))ll s sc
(23)
Theorem C.2. Define St, 0g be the states and actions resulting froul current policy using ground- truth dynamics P and reward function r and similarly define 8 ag using learned functions Assume the learned reward function f to be Lip-cips and the estimated optimal Q-function Q* to be Li-Lipshhiz. Assume the estimation errors of learned functions are bounded as in Assumption C.1 Then we have the error between barch-based Q-value estimation f(so, ag) and ground-truth Q-value Q* (Su, Go) bounded as:
functioos P and ft
< 1- 7
4r + ( 7
[s(o, ao) - (* a0)l
fy + =500 e, "
(24)
1 g
1--
16

---

Published as a conference paper at ICLR 2025
Proof.
(I up) - Qx0
H
7 71
7 7 an) +7" max Q (8 an)
- E
Σ
max Q" (s an) (r(ss,, ar) +7" an
0
H
< E
F(s0, ao) r(s0, a0)
+
a)) ((0) + +y
-
max AN
(5 an) max Q" (8H, aH)
1=1
(s) ao) r(50. ao)
Σ
? (()) r(se, ac) + +
max Q* aH) max Q" (sH 0H)
ax
ax
t=1
H-1
Se,
Σ
-((
an)
+
Q* (8 max oz
max og
Q" (sH, an))
(25)
For the second term in inequality
25
f(s a1) r r(st, ai)
|(s a;) f(st, ai) + f(s), 0s) - r(s( r(s),a)) ail
Moii ai) -f(ss, + [|((se,))) r(s),
< br s - 51 +
<I, + cr
(26)
For the third term in inequality arg maxam (sH, an). Then we have:
25 let al = arg max (" an) and ai =
E
max (s max Q" (sH
ax
ax
E
(' d Q* a/r)
E
Q'(5 dH/)
(s a Q" (sH, ('H
s Q (s a11
(* o ||| + Q azz)
<LQ
sH ss
+ cc
SLoE, + EQ
(27)
Substitute inequalities
26 and 27 into 25
[f ao) - Q" (s0.
H-1
Se,
(Lres + ter) ( (Loes +EQ)
M
Er + 1 1 er + =0 6s + "
(28)
L
Analysis. We expect the search-based Q-values f(8 a) have an upper error bound no greater than the estimated Q-values Q' (s, a), which is formulated as the following inequality:
1- 7 1- 6r +
7 --
1-7 4r + (m((co c0 6 S <<o
(29)
17

---

Published as a conference paper at ICLR 2025
Based inequality 29 derive the following condition:
1
7 L, - 7 L, LQ 701
+ (1-
1 S EQ
(30)
The inequality (30) means that if the weighted of rewards estimation enror &, and state transition estimation error 6, are less than or equal to the Q-values estimation error <q, then the search-based optimal Q-values have lower upper-bound of error than estimated optimal Q-values.
D GAMES
We select 20 Atari games maintaining the difficulty distribution of full Atari 2600 games defined by Gulcehre et al. (2020), which includes 9 easy games, 9 medium games, and 2 hard games. We use 15 out of 20 games for training and the remaining 5 for OOD generalization experi- ments. The 15 training games are: Phoenix, Centipede, SpaceInvaders Carnival, NameThisGamt Assault, Atlantis, DemonAttack, BeamRi der, ChopperComma Seaquest, TimePilot, StarGunner, Berzerk, Zaxxon. The 5 held-out games are: Pong, Robotank, YarsRevenge, Gravitar, MsPacman. Details about the size of action spaces and game difficulties are shown in Table
Table 7: Atari Games: Name, Game difficulty, Action Space, and Type.
Game
Difficulty
Action Space
Type
Assault
Medium
Train
Atlantis
Hard
18
Train
BeamRider
Medium
Train
Berzerk
Hard
Train
Carnival
Medium
Train
Centipede
Medium
Train
ChopperCommann
Easy
18
Train
DemonAttack
Easy
Train
Gravitar
Easy
Fine-tune
MsPacman
Medium
Fine-tune
Name Game
Easy
6
Train
Phoenix
Easy
Train
Pong
Medium
Fine-tune
Robotank
Medium
Fine-tune
Seaquest
Easy
Train
SpaceInvaders
Easy
Train
StarGunner
Medium
18
Train
TimePilot
Easy
Train
YarsRevenge
Medium
Fine-tune
Zaxxon
Easy
18
Train
E EXPERIMENTAL DETAILS
E.1 IMPLEMENT DETAILS
E.I. JOWA
We implement JOWA based on the codes of IRIS (Micheli et al. 2022 We train the tokenizer
VQ- VQ-VAE using the following loss function:
C(E, = | D(x) + [}s((((z)) - E(=) [|sg(E(z)) - E(x) + Cperceptued D()))
where F,, D, encoder, decoder, and embedding table respectively. sg(.) is the stop-gradient operator. The last term is the perceptual loss (Johnson et al.] 2016). We list the hyperparameters of VQ-VAE Table and After the first stage of pretraining, the VQ-VAE is frozen.
18

---

Published as a conference paper at ICLR 2025
Table 8: Encoder / Decoder hyperparameters. We list the hyperparameters for the encoder, the same ones apply for the decoder.
Table 9: Embedding table hyperparameters.
Hyperparameter
Value
Hyperparameter
Value
Frame dimensions (h, w)
84 84
Vocabulary size
2048
Layers
Tokens frame (K)
36
Residual blocks per layer
Token embedding dimension
512
Channels in convolutions
64
Self-attention layers at resolution
6 12
In addition to the vocabulary embedding and position embedding, we add learnable task embed- ding for observation tokens and action tokens respectively. Our transformer are based minGPT with Fltentin (Dao et al. 2022) for acceleration. The hyperparameters JOWA's transformer
backbone listed in Table 10 and
Table 11:
Different hyperparameters trans-
Table 10: Same hyperparameters of transformer for 3 JOWA variants.
former for 3 JOWA variants.
Model
Layers
Hidden size
Heads
Hyperparameter
Value
JOWA-40M
4
512
8
max sequence tokens
296
JOWA-70M
6
768
12
dropout rate
0.1
JOWA-150M
12
768
12
Table 12: Hyperparameters of Q-heads for JOWA variants.
Model
Layers
MLP Hidden dimension
Number heads
Dropout
JOWA 40M
768
0,01
JOWA-70M
1024
JOWA-150M
1792
3
0.01
The observation predictor, reward predictor, and terminal predictor are 2-layers MLP. The Q-heads
are MLP with dropout, layer normalization, and Mish activations from Hansen et al. (2024). The hyperparameters of Q-heads for JOWA are shown in Table 12 The training hyperparameters of JOWA are shown Table 13
1.1 MTBC
We implement MTBC based on JOWA. We remove the observation predictor, reward predictor, and terminal predictor. We change the output dimension of Q-heads to 18 and train the heads 18-class classification problem. All hyperparameters are kept the same as JOWA.
E. EDT
We use the official code for EDT We implement EDT-200M based on the architecture configu- ration MGIT-220M,, which is shown in Table 14 We change the batch size 512 and keep other hyperparameters the same as its original configuration. We enable data augmentation (random cropping and random rotation) for EDT.
E.1.4 MGDT
We implement MGDT based on the codes of EDT. We use {ou+e, Re++, as+is T as the input sequences, remove the expectile regression loss Lmax and observation prediction loss Lobvervation and
19

---

Published as a conference paper at ICLR 2025
Table 13: Training hyperparameters of JOWA.
Hyperparameter
Value
Optiizzrr (VQ-VAE)
Adam
Optimizer (except VQ-VAE)
Adam W
Learning rate (VQ-VAE)
0.0001
Learning rate (except VQ- VAE, stage 1)
0.0001
Learning rate (except VQ-VAE, stage
0.00005
Batch size (VQ-VAE)
2048
Batch size (except VQ-VAE)
512
Weight decay (except VQ-VAE)
Gradient clip
1.0
Discount factor (7)
0.99
Target Q update frequency
1000
Distributional Q
[-10, 301
Number atoms
Coefficient of CQL (a)
0.1
Coefficient of Lworss (3)
0.1
Table 14: Hyperparameters of transformer for 2 MGDT variants.
Model
Layers
Hidden size
Heads
MGDT-40M
6
768
12
MGDT-200M
10
1280
20
add the reward prediction loss rewrite the codes EDT into MGDT. We enable data augmentation (random cropping and random rotation) for MGDT. The hyperparameters of transformer for two MGDT variants are listed Table
E. SCALED-QL
We implement pytorch version of Scaled- QL from scratch, referring to the jax version of its official preliminary codesf] We use ResNet-101 the representation backbone, followed by 3-layers MLP with 1024 hidden neurons and an output layer. We replace the batch normalization in ResNet with group normalization and use learnable spatial embeddings to aggregate the outputs of the ResNet instead of global mean pooling. Before the output layer, we normalize the feature e as Tel The training hyperparameters of Scaled-QL are listed in Table 15
For fair comparison, all methods trained with the same batch size of 512 for 1.75M gradient steps. For reporting results, we report the performance of the agent at the end of pretraining.
E.2 FINE-TUNING PROTOCOL
We uniformly draw 5k transitions from expert-level DQN-Replay (Agarwal et al. 2020) (last 20% of the original dataset) for each held-out game. Each game was fine-tuned separately to model's 's transfer performance for a fixed game. we fine-tuned all methods using a batch size of 32 and learning rate of 0.00005 for 50k gradient steps. For reporting results, report the performance of the agent snapshot that obtain the highest score during fine-tuning.
the
For JOWA in the second fine-tuning stage, we both the planning horizon and the beam width to 2 for all fine-tuning experiments. Thus we sample batch of 6-steps segments, using planning algorithm to synthesis the last steps. For other baselines, we enable random cropping and random rotation for data augmentation.
20

---

Published as a conference paper at ICLR 2025
Table 15: Training hyperparameters Scaled-QL.
Table 16: Evaluation settings of Atari.
Hyperparameter
Value
Hyperparameter
Value
Optimizer
Adarn
Learning rate
0.0002
Sticky actions
No
Batch size
512
Grey-scaling
Truc (84, 84)
Gradient clip
1.0
Observation down-sampling
Discount factor (?)
0,99
Frames stacked
4
Target Q update frequency
2000
Frame skip (Action repetitions)
4
Distributional Q
[-20, 20]
Terminal condition
Game Over 108K 0.001
Number atoms
51
frames per episode
Coefficient of CQL (a)
0.05
Evaluation noise Coal
n-step returns
3
E.3 EVALUATION PROTOCOL,
For all methods, each game score is calculated by averaging over 16 model rollout episode trials. To reduce inter-trial variability, we do not use sticky actions during evaluation following Lee et al. (2022); Kumar et al, (2023). Following standard protocols on Atari, we evaluate a noised version of the policy with an apsilon-greedh scheme, with Eeval 0.001. The evaluation settings of Atari are shown in Table [6
For the expert action inference MGDT and EDT, we set the inverse temperature k to 10. For the planning of JOWA, set the planning horizon H to 2 for all games. The beam width are set according to the size valid action space of each game. Specifically, we set beam width to 2 if the valid action space size is less than 10, otherwise set K in {3, 4}. The planning hyperparameters for each game are shown Table
Table 17: Planning hyperparameters of JOWA during evaluation.
Game
planning horizon
beam width
Action space
Assault
2
7
Atlantis
2
18
BeamRider
2
9
Berzerk
2
4
18
Carnival
2
6
Centipede
2
18
ChopperCommand
2
18
DemonAttack
2
2
6
NhimaaaG
2
6
Phoenix
2
8
Seaquest
18
SpaceInvaders
2
2
6
StarGunner
2
18
TimePilot
2
10
Zaxxon
18
F MORE EXPERIMENTS AND RAW SCORES
We summarize the raw of fine-tume fine-tning experiments, ablation studies on planning algorithms and training choices in Table 18 19 and 20 respectively.
Details of the comparison of planning algorithms. python rather than C++ for fair speed comparison. Note that Muzero-style (Schrittwieser et al. 2020) MCTS requires access to the V-function and policy networks x while JOWA only estimates the optimal Q Qufuctioo. Therefore, conduct grid search the following choices for MCTS:
implement both planning algorithms in
1. Compute the V-value using Q. mean () or m.aax ()
21

---

Published as a conference paper at ICLR 2025
Table 18: Offline fine-tuning performance on unseen games using 5k transitions, measured in terms of DON-normalized score, following Lee et al. ) Kumar et al. (2023).
MTBC 120M
MGDT
EDT 200M
SQL 80M
JOWA
JOWA-150M
Game
Random
DQN
200M
150M
(scratch)
Gravitar
173.0
473.0
35.7
250.0
253.3
137.5
273.3
83.3
MsPacman
307.3
3085.6
905.0
1290.3
1210.7
1040.2
2016.7
786.7
Pong
-20.7
19.5
5.8
9.7
11.3
13.7
17.7
8.8
Robotank
2.2
63.9
6.8
16.0
15.5
19.7
25.0
11.0
YarsRevenge
3092.9
18089.9
7987.5
10886.3
11276.9
10838.5
17506.2
6507.0
Mean
0.000
1.000
0.164
0.422
0.430
0.360
0.647
0.196
Median
0.000
1.000
0.215
0.354
0.325
0.284
0.615
0.173
IQM
0.000
1.000
0.205
0.377
0.380
0.355
0.647
0.181
Table 19: Raw scores on games for JOWA-150M evaluated with different planning algorithms.
Game
Random
Human
w/o planning
MCTS
Ours
BeamRider
363,9
16926.5
864.5
1137.0
3498
Berzerk
123.7
2630.4
396.9
440.0
739
Carnival
0
3800
5560.0
3340.0
5316
ChopperCommand
811.0
7387.8
806.2
1850.0
3812.5
Seaquest
68.4
42054.7
267.5
760.0
2725
TimePilot
3568
5229.2
662.5
4100.0
3669
Zaxxon
32.5
9173.3
12.5
50.0
2163
#Mean HNS
0.000
1.000
-0.02
0.221
0.378
IQM HNS
0.000
1.000
0.03
0.134
0.237
2. We employ an energy-based policy to compute action probability, ie., T-|ss) = (ssnnax((((ss)))))) where is the temperature. We search in (0.01, 0.1, 0.5, 0.7, 0.8, 0.9, 2, 3, 5, 10].
3. Use the action of most visited or most valuable children the root state as the optimal action.
4. Search the max depth of the tree in [1,
Finally, we use the configuration with V=Q max x (), t 0.9, and action of most valuable children as the optimal action for all games while searching H in (1, 2, 4, 6} for each game. Even though we believe that we conduct sufficiently adequate hyperparameter search, MCTS still performs worse than ours, as shown in Table 19
We compare these two algorithms performs poorly. The results show that our planning algorithm exceeds MCTS by 71.0% Mean HNS. Moreover, we find that MCTS highly sensitive to the temperature and max depth H, and inappropriate values of hyperparameters can even degenerate the policy into a randomized policy. However, its long execution time makes inconvenient to tune hyperparameters. Because of this, we have not yet found hyperparameters that can make MCTS perform reasonably on the other 8 pretrained games. We will include MCTS as an optional planning algorithm in our open-source evaluation code.
games where bare JOWA-150M (i.e., without planning)
Ensemble of Q-value heads. We run JOWA-40M with ensembled Q-heads, where the return dis- tribution is weighted sum of distributions from multiple Q-heads. The original configuration of JOWA-40M uses a single Q-head. In this experiment, we set the number of Q-heads to and com- pare 2 ensemble approaches: equal weights random weights like REM (Agarwal et al.) 2020). We report the results in the forth main column of Table 20 Although the equal weighted ensem- ble slightly outperforms random weights, we have not observed significant improvements of multi Qehaads over single Q-head. We leave the better way of distributional-O ensemble for future work.
Effects of task embedding. We run JOWA-40M without task embedding, using only the sum of vocabulary embedding and position embedding as input for transformer. We report results in the
22

---

Published as a conference paper at ICLR 2025
Table 20: Raw scores on the Atari games for various training choices. The mean, median, and IQM human-normalle score are shown in the last 3 rows and the best scores are markded in bold.
Different training losses
Q-heads ensemble
No task embedding
Synthetic data
Game
Origin
CQL
No Leail
sg (astoon
MSE
Equal
Random
in pretraining
Assault
1423.5
857.9
452.7
1628.8
1258.6
1428.5
655.7
Carnival
5560
4144.6
5560
2120
620
5154.3
4160
5974.2
3492,9
Centipede
5018.9
1494.1
8146
$568.7
4097.3
5592.5
6450
3725.4
3253
NameTTisGGme
12208.1
9307.1
4407.3
2108.8
1315
12148.6
9420
7700
5080
Phoenix
4740
140
2036.7
1920
1190
4610
4941.7
4020
193.3
Spucelnvaders
1201.7
605
323.4
539.3
260
958.4
1283.3
575.4
786.3
Mean HNS
1.183
0.613
0.917
0.293
0.189
1.209
1,026
0.964
0.448
Median HNS
1,078
0.696
0,489
0.304
0,118
0.975
0.921
0.721
0.452
IQM HNS
1.123
0.637
0.659
0.307
0,126
1.049
0.931
0.824
0.464
Table 21: Offline fine-tuning performance JOWA-150M on unseen games with various data qual-
ity, measured in terms of DON- score, following Lee et al. (2022
Kumar et al. (2023).
Game
Random
DQN
Expert
Suboptimal
Highly-suboptinal
Gravitar
173.0
473.0
273.3
317.8
296.0
MsPacman
307.3
3085.6
2016.7
1005.5
1126.8
Pong
-20.7
19.5
17.7
13.2
13.8
Robotank
2.2
63.9
25.0
14.6
8.5
YarsRevenge
3092.9
18089.9
17506.2
15085.0
9755.4
Mean
0.000
1.000
0.647
0.516
0.422
Median
0,000
1.000
0,615
0,483
0.410
IQM
0.000
1.000
0.647
0.511
0.383
fifth main column of Table 20 Observe that the addition of task embedding improves 21.9% and 35.7% for mean and median human-normalln normalized score respectively.
Comparing with the-tuning using expert-level data, we addi- tionally fine-tune JOWA- 150M using 5k suboptimal and highly-suboptina transitions. Specifically, the suboptimal and the the initial 20% of the DON-Replay dataset, respectively. The results Table 2T show that the fine- tuning performance strongly correlates with data quality. The mean RQN-normalized scores for expert, suboptimal, and highly-suboptimal data are 0.647, 0.516, and 0.422 respectively.
Fine-tuning with non-expert data.
transitions are uniformly sampled from the complete and
Discussion of emergent behaviors. Capability emergence refers to the phenomenon where large models suddenly exhibit after reaching certain scale thresholds in model size. For example, JOWA-150M achieves a sig- nificantly higher score than JOWA-70M and other baselines on Zaxxon, which seems like the "emergent" phenomenon. However, after examining the reward of Zaxxon, we conclude that this "emergent behavior" is actually attributable to the nonlinear reward function. In addition the com- mon rewards of around 100, Zaxxon also defines huge reward of nearly 5000. JOWA-70M gets this huge reward in of the 16 rollouts while JOWA-150M gets it in 7 rollouts during evaluation. When we clip the rewards to the range JOWA--0NN and JOWA-150M achieve scores of 0.08 and 0.44 respectively, indicating no true "emergent" phenomenon. [Schaeffer al. (2024) draws similar conclusion that "nonlinear discontinuous metries cause the emergent abilities".
abilities significantly enhanced performance, typically occurring
23


Reviews:
[{"review_id": "drePSpW640", "reviewer": "Reviewer_zGRd", "summary": "This paper presents JOWA, an offline model-based reinforcement learning method aimed at scaling and generalizing multi-task RL through a shared transformer-based architecture. JOWA is designed to stabilize temporal difference learning for large models by integrating world modeling with Q-value criticism, thus leveraging a shared transformer backbone for both tasks. The paper introduces a novel parallelizable planning algorithm to counter Q-value estimation errors, achieving more consistent policy identification during inference. The model is pre-trained on Atari games with minimal data and reportedly outperforms existing baselines, demonstrating some generalization capacity to new tasks.", "questions": "1. How might JOWA perform in a single-task setting compared to standard model-based methods? Would the shared architecture yield efficiency or performance gains? If it is worse than training from scratch with model-based RL, why the Atari games are good benchmarks for evaluating the scaling laws?\n\n2. The author mentioned the beam search, so why not use MCTS, which is proven effective in MuZero and EfficientZero. And the author didn't compare with other model-based RL algorithms, such as Dreamer.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. Scalability: JOWA’s design showcases robust scaling potential, as performance improves with model size without the usual TD-learning instability issues.\n\n2. Detailed Ablation Studies: The authors conducted extensive ablations, examining the impact of core design elements such as task embeddings, training losses, and synthetic data usage.", "weaknesses": "1. Missing related work and explanations for the architecture of the world-action modeling. The work proposes to use VQ-VAE for the representation learning in Atari games, while there is a work named Forward-Inverse Cycle Consistency (FICC), which uses VQ-VAE for the offline Atari dataset to learn representations and action embeddings. It seems this pipeline is so similar to the FICC, however, the author didn't mention the difference between JOWA and FICC. (Both use offline Atari datasets for model-based RL and learning universal policy under VQ-VAE loss terms)\n\nYe, W., Zhang, Y., Abbeel, P., & Gao, Y. (2022). Become a proficient player with limited data through watching pure videos. In The Eleventh International Conference on Learning Representations.\n\n2. Overclaim of the generalization ability. The experiments are based on Atari games, while using offline atari dataset for pre-training. Extending this approach to more complex control task scenarios (e.g., robotics or high-dimensional continuous control tasks) remains unproven. I'm not sure about the effectiveness of the scaling results.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "EDA9RTKJrR", "reviewer": "Reviewer_7Vxq", "summary": "This paper introduces a model-based RL agent, JOWA, which employs a transformer architecture and jointly optimizes world dynamics and Q-values across different environments. To compensate for inaccurate Q-values, the learned world model enables JOWA to search out the optimal policy via planning during inference time. Empirically, JOWA demonstrates impressive performance in a low-data regime by training on 15 Atari games. Moreover, JOWA's performance scales up with model sizes.", "questions": "How's the model performance when you further scale up the model size, e.g., to 1B?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The proposed JOWA outperforms existing SOTA methods by a large margin. The perform scales up with model sizes\n\n- The joint optimization of world and action models stabilizes large-scale multi-task offline RL training.\n\n- The ablation studies comprehensively study the key design choices of the proposed methods.", "weaknesses": "- The proposed method combines the best offline RL training techniques, leveraging the world modeling loss to stabilize Q-value learning. The empirical performance is impressive. However, the technical novelty is thus limited.\n\n- By taking a closer look at Table 2, we can see that the 150M variant does not consistently outperform the 40M and 70M variants on all tasks. For example, the 40M variant achieves the highest score on Centipede, while the 70M variant excels on NameThisGame, SpaceInvaders, and StarGunner. Can you explain the phenomenon a bit more? Is it because of the training instability? Can further increase the model size and training iterations solve the issue?\n\n- It would be insightful to examine the emergent behaviors that arise when scaling up model size. For example, the 150M JOWA achieves a significantly higher score than the 70M JOWA and the other baseline methods on *Zaxxon*. Does the 150M JOWA agent exhibit any distinct emergent behaviors that could explain this performance improvement?\n\n- This work misses the citation for [1,2]\n\n[1] Li et al., Multi-task batch reinforcement learning with metric learning. NeurIPS 2020.\n\n[2] Li et al., Offline reinforcement learning with closed-form policy improvement operators. ICML 2023", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "jaa5k4Erj2", "reviewer": "Reviewer_pKew", "summary": "This paper introduces JOWA, a Jointly-Optimized World-Action model for offline RL. By optimizing a shared transformer backbone for world-action modeling and employing an efficient planning algorithm to improve policy search, JOWA achieves good performance on pretrained games with limited data and demonstrates superior generalization to new games with minimal fine-tuning trajectories.", "questions": "Please see the weaknesses section.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper is clearly written and easy to follow.\n2. This paper presents sufficient experimental results to demonstrate the validity of its proposed method.", "weaknesses": "1. A large generalist TD-MPC2 agent is capable of performing a variety of tasks across multiple domains. I wonder if the proposed method is better than TD-MPC2 in the offline setup.\n2. Extending the experiments beyond Atari to more complex environments like Kitchen or Meta-World would offer stronger validation of the proposed method's effectiveness.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "kHtAWLw0yx", "reviewer": "Reviewer_pLGR", "summary": "The paper introduces JOWA (Jointly-Optimized World-Action model), a novel offline model-based reinforcement learning (RL) agent designed for generalization across multiple tasks, specifically focusing on scaling RL with image-based observations. JOWA is pretrained on multiple Atari games using a shared transformer backbone that jointly optimizes a world model and action model. This setup allows the model to learn both general-purpose representations and decision-making skills.\n\nJOWA stabilizes temporal difference learning by incorporating world modeling as a regularizer, enabling large models to scale effectively. The proposed framework also features a parallelizable planning algorithm to improve Q-value estimation and policy search at inference time. Experiments show that JOWA outperforms state-of-the-art offline RL methods, achieving 78.9% human-level performance on pretrained games and demonstrating strong sample-efficient transfer to novel games with minimal fine-tuning.", "questions": "- Why was the sequence length relative low (8)?\n- Could JOWA learn useful policy on new games using non-expert transitions?\n- It seems the number for MGDT are very different from the original paper (MGDT achieved 93% IQM HNS on 40M model). Why is this?\n- Is the action-part module (CQL) hard to tune in this method?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- Extension from Sequential Modeling: JOWA effectively adapts transformer-based architectures from NLP to reinforcement learning, leveraging tokenization for image observations and combining offline RL with online planning to scale across tasks.\n- Strong Baseline Performance: It outperforms state-of-the-art methods, showing significant improvements on Atari benchmarks while requiring only 10% of the dataset.\n- Adaptability: JOWA demonstrates efficient transfer to new tasks with minimal fine-tuning, highlighting its robustness and generalization capabilities.", "weaknesses": "- Limited Originality: While JOWA effectively combines existing state-of-the-art techniques, it lacks substantial novelty in introducing new ideas. The core contributions largely build upon well-established methods, such as transformers, tokenization, and offline RL.\n- Reliance on Expert Data for New Games: While JOWA demonstrates strong performance in fine-tuning to new tasks, it heavily relies on expert-level data for unseen games. This dependence on high-quality data might limit its applicability in scenarios where such data is not readily available or costly to obtain.\n- Excessive Discretization: The paper discretizes the reward into the set {-1, 0, 1}, which may lead to a loss of important nuance and granularity in reward signals. This overly simplified reward structure might not capture the full complexity of more sophisticated environments, limiting the model’s performance in tasks requiring finer reward distinctions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "PtX2IAR3iI", "reviewer": "Reviewer_XU9g", "summary": "The paper focuses on the important research question: \"Can image observation-based world model scale offline RL across multiple tasks while enhancing generalization to diverse unseen tasks?\". The paper presents JOWA, which 1) jointly trains a Transformer model that includes world model, policy, critic using the offline data, 2) planning with beam-search in inference-time using the trained world model. Experimental results in Atari tasks shows that JOWA trained with 10% of subsampled 15 tasks significantly outperforms previous multi-task offline RL approaches not only in in-distribution tasks, but also in fine-tuning with 5k expert demonstrations. Moreover, JOWA shows favorable scaling property respect to the number of parameters. Exhaustive ablation studies show the importance of individal components of JOWA.", "questions": "Q1. Can Theorem 1 be applied for planning of JOWA?\n* Theorem 1 shows that error caused by model-based planning can be bounded. However, it assumes that the critic $Q(s,a)$ is the estimate of the optimal Q-function (i.e. $\\mathbb{E}[\\sum_t \\gamma^t r_t]$). However, the critic is learned via CQL, which tends to underestimate the Q-values, as discussed in [1]. Thus, $Q(s, a) - Q*(s, a)$ will not be bounded by the term mentioned in Theorem 1; Q(s, a) will be much smaller than $Q^{*}(s, a) = \\mathbb{E}[\\sum_t \\gamma^t r_t]$ in reality.\n* Can you share your thoughts on this? Please correct me in the point if I'm wrong.\n\nQ2. Scaling property in fine-tuning tasks\n* Can you provide the fine-tuning results for JOWA-40M, JOWA-70M?\n* I think it will be helpful on investigating scaling effects in fine-tuning.\n\nQ3. Extending JOWA to continuous environments\n* It would be exciting if JOWA can be also extended to continuous tasks. For example, continous tasks like DMControl are relatively simple but challenging.\n* However, I understand that the training process is time-consuming and authors should put a lot of effort on this. Thus, I'm not strongly suggesting (but welcome) the authors to answer this question via additional experiments; if then, telling your thoughts or ideas on how to extend this work to continuous environments will be helpful for potential users of JOWA and future works.\n\n[1] Mitsuhiko Nakamoto et al., Cal-QL: Calibrated Offline RL Pre-Training for Efficient Online Fine-Tuning, NeurIPS 2023.", "limitations": "", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "S1. JOWA significantly outperforms previous multi-task offline RL approaches.\n* Experimental results of Table 2 shows that JOWA gives state-of-the art result in multi-task offline RL with 10% subsampled Atari (Mean 0.456, IQM 0.789)\n* Experimental results of Table 3 shows that JOWA significantly outperforms prior works in fine-tuning (Mean 0.647, IQM 0.647)\n\nS2. JOWA shows favorable scaling property.\n* Increase in the number of parameters leads to better performance.\n* Compared to other methods, the performance increase is significant.\n\nS3. Ablation studies shows useful insights on JOWA's performance.\n* Planning with world model largely contributes the performance.\n* Training policy jointly with the world model is important for multi-task agents.\n\nS4. Open-source codes & checkpoints.\n* Those resources will be helpful for the community.", "weaknesses": "W1. Evaluated only in Atari\n* For example, continous tasks like DMControl are relatively simple but challenging.\n* It will be exciting to see if JOWA can be also applied in continuous tasks.\n\nW2. Discrepancy between Theorem 1 and planning of JOWA\n* Although the planning emprically improves the performance, it seems that the theory slightly differs with the planning of JOWA. Specifcally, Theorem 1 assumes critic is the estimate of the optimal Q-function, but JOWA uses a conservative estimate (CQL). Thus, $Q(s, a) - Q*(s, a)$ will not be bounded by the term mentioned in Theorem 1; Q(s, a) will be much smaller than Q*(s, a) in reality.\n* Please correct me in the point if I'm wrong.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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:

Decision: Accept (Poster)