Scaling Offline Model-Based RL via Jointly-Optimized World-Action Model Pretraining

## Abstract

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: J ointly- O ptimized W orld- A ction 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.

## 1 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-modality. Notable examples include large language models (Brown et al., ; Ouyang et al., ; Touvron et al., ; Team et al., ) and large vision-language models (Zhu et al., ; Liu et al., ; Bai et al., ), which deliver outstanding performance across a wide range of tasks and adapt quickly to new ones through few-shot or in-context learning. Their success is largely driven by the scaling law (Kaplan et al., ), 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., ; Ota et al., ; Sokar et al., ), explaining the continued dominance of shallow CNN networks in vision-based RL tasks (Mnih, ). While some efforts have been made to scale offline RL across multiple tasks (Lee et al., ; Reed et al., ; Xu et al., ; Wu et al., ), they predominantly rely on supervised learning (SL) approaches, such as conditional behavior cloning, rather than temporal difference (TD) learning, and heavily rely on large amounts of expert trajectories. Kumar et al. () 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, scaling TD-based offline RL for simultaneous general-purpose representation and decision-making remains a critical challenge. While Hansen et al. () 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 a distinct research area within computer vision and multi-modality, primarily focusing on conditional video generation (Hong et al., ; Blattmann et al., ; Yu et al., ; Yang et al., ; Bruce et al., ). Notably, SORA (Brooks et al., ) 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-based world model scale offline RL across multiple tasks while enhancing generalization to diverse unseen tasks?”

To address this question, we introduce JOWA: Jointly-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-propagated from both world modeling and TD losses, enabling joint optimization. The world modeling loss acts as a regularizer, stabilizing TD-learning for large models. Additionally, we propose a provably efficient and parallelizable planning algorithm to compensate 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 a single model to play 15 Atari games, similar to Lee et al. (); Kumar et al. () 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 diverse nature of games, each with unique dynamics, rewards, visuals, and agent embodiments. To further test JOWA’s generalization ability, 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 capable of training a single high-performing generalist agent across multiple Atari games. JOWA attains 78.9% human-level performance on pretrained games using only 10% of the original dataset (Agarwal et al., ), outperforming existing state-of-the-art large-scale offline RL baselines by 31.6% on averange. Second, we demonstrate that JOWA unlocks scaling trends, with performance improving as model capacity increases. Third, JOWA enables sample-efficient transfer to diverse unseen games with 64.7% DQN-normalized score using only 5k transitions per game, surpassing baselines by 34.7% on average. Our ablation studies highlight the significance of two key design features of JOWA: joint optimization and planning, along with other training choices. Due to the extensive training time, we will release all training, evaluation, and fine-tuning codes, as well as model weights, to support future research.

## 2 Related work

##### Offline Reinforcement Learning.

Offline RL algorithms learn a policy entirely from the static offline dataset without online interactions. Model-free offline RL incorporates conservatism to mitigate extrapolation error (Jin et al., ) primarily through policy constraints (Fujimoto et al., ; Kumar et al., ; Nair et al., ; Kostrikov et al., ; Fujimoto & Gu, ) and value regularization (Kumar et al., ). In contrast, model-based offline RL approximates the environment using learned world models and performs conservative policy optimization (Lu et al., ; Yu et al., ; ; Schrittwieser et al., ). While these works focus on single-task settings, our work explores scaling offline model-based RL across diverse, challenging multi-task Atari games (Lee et al., ; Kumar et al., ; Wu et al., ) aiming for sample-efficient transfer to novel games. We employ CQL (Kumar et al., ) and COMBO (Yu et al., ) to estimate conservative Q-values, leveraging their simplicity and effectiveness on vision-based offline RL tasks.

##### Mutli-Task Reinforcement Learning.

Multi-task reinforcement learning (MTRL) aims to learn a shared policy for diverse tasks, with various approaches proposed in the literature (Teh et al., ; Song et al., ; Xu et al., ; Yang et al., ; Sodhani et al., ). A common approach is to formulate the multi-task model as task-condition (Yu et al., ), such as language-conditioned tasks (Lynch & Sermanet, ; Ahn et al., ; Jang et al., ) and goal-conditional RL (Plappert et al., ). In multi-task offline RL, conditional sequence modeling approaches based on decision transformer (Chen et al., ) or diffusion model (Janner et al., ) typically rely on large amounts of expert trajectories (Reed et al., ; Xu et al., ; Lee et al., ; Wu et al., ; Hu et al., ; He et al., ).
Beyond conditional behavior cloning (BC) methods, Scaled-QL (Kumar et al., ) 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 jointly-optimized world-action model. Table 1 compares experimental environments and open-source status of multi-task offline RL algorithms. Following Lee et al. (); Kumar et al. (); Wu et al. (), we focus on the multi-game regime, which presents greater challenges due to high-dimensional observations and diverse, stochastic environment dynamics. Due to the prohibitive long training time, we advocate for open-sourcing codes and checkpoints to facilitate progress in this field.

**Table 1: Comparison of methods in multi-task offline RL. $\clubsuit$ and $\spadesuit$ represent two training paradigms of agents, conditional BC and TD-learning, respectively.**
| Method | Experimental environment | Open source |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Benchmark | Observation | Action | Dynamics | Train | Eval | Check- |  |  |
| space | space | per task | across tasks | points |  |  |  |  |
| MTDIFF$\clubsuit$ (He et al., ) | Meta-World | state | continous | deterministic | same or | ✓ | ✓ | ✗ |
| HarmoDT$\clubsuit$ (Hu et al., ) | or DMControl | similar | ✓ | ✓ | ✗ |  |  |  |
| TD-MPC2$\spadesuit$ (Hansen et al., ) | ✓ | ✓ | ✓ |  |  |  |  |  |
| MGDT$\clubsuit$ (Lee et al., ) | Atari 2600 | pixels | discrete | stochastic | diverse | ✗ | ✓ | ✓ |
| Elastic DT$\clubsuit$ (Wu et al., ) | ✓ | ✓ | ✗ |  |  |  |  |  |
| Scaled-QL$\spadesuit$ (Kumar et al., )(^1^11Scaled-QL released the preliminary code in [https://tinyurl.com/scaled-ql-code](https://tinyurl.com/scaled-ql-code), which is not run-able out-of-the-box.) | ✓<br>✗ | ✓<br>✗ | ✗ |  |  |  |  |  |
| JOWA (ours)$\spadesuit$ | ✓ | ✓ | ✓ |  |  |  |  |  |

## 3 Preliminaries and problem setup

### 3.1 Online distributional rl (C51)

In distributional RL, the distribution $Z$ over returns replaces the Q-value in the Bellman optimality equation. The Q-value is the mean of the value distribution $Z$ that can be computed through a distributional Bellman optimality operator (Bellemare et al., ),

$$ $\mathcal{T}^{*}Z(s,a):\overset{D}{=}R(s,a)+\gamma Z(s^{\prime},\arg\max_{a^{ \prime}}Q(s^{\prime},a^{\prime}))$ (1) $$

where the formula $Y:\overset{D}{=}U$ denotes equality of probability laws, that is the random variable $Y$ is distributed according to the same law as $U$. The C51 algorithm (Bellemare et al., ) models $Z(s,a)$ using a discrete distribution supported on $N$ fixed atoms $z_{1}\leq\cdots\leq z_{N}$ uniformly spaced over a predetermined interval. Given a current value distribution, C51 applies a projection step to map the target distribution onto its finite element support and optimizes as follows:

$$ $\mathcal{L}_{\text{TD}}=\displaystyle D_{\mathrm{KL}}(\mathcal{T}^{*}Z_{\theta ^{-}}(s,a)\|Z_{\theta}(s,a))$ (2) $$

### 3.2 Value regularization based Offline rl (CQL and COMBO)

To be conservatism on unseen actions, CQL (Kumar et al., ) 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:

$$ $\mathcal{L}_{\text{CQL}}=\alpha\left(\mathbb{E}_{s\sim\mathcal{D}}\left[\log \left(\sum_{a^{\prime}}\exp{\left(Q_{\theta}(s,a^{\prime})\right)}\right) \right]-\mathbb{E}_{s,a\sim\mathcal{D}}\left[Q_{\theta}(s,a)\right]\right)+ \mathcal{L}_{\text{TD}}$ (3) $$

COMBO (Yu et al., ), a model-based variant of CQL, uses the following loss function:

$$ $\mathcal{L}_{\text{COMBO}}=\alpha\left(\mathbb{E}_{s,a\sim\rho_{\hat{\mathcal{ M}}}(s,a)}\left[Q_{\theta}(s,a)\right]-\mathbb{E}_{s,a\sim\mathcal{D}}\left[Q_ {\theta}(s,a)\right]\right)+\mathbb{E}_{s,a,s^{\prime}\sim d_{f}}\left[ \mathcal{L}_{\text{TD}}\right]$ (4) $$

where $\hat{\mathcal{M}}$ is the imagined Markovian Decision Process induced by the world model. $\rho_{\hat{\mathcal{M}}}(s,a)$ is the occupancy measure of current policy within $\hat{\mathcal{M}}$. $d_{f}$ is an $f$-interpolation between the offline dataset and synthetic rollouts from the world model: $d_{f}=fd_{\mathcal{D}}+(1-f)d_{\hat{\mathcal{M}}}$, where $f\in[0,1]$.

### 3.3 Problem setup

We consider a multi-task offline RL problem: given a static dataset of transitions $\mathcal{D}=\{(s_{t},a_{t},r_{t},d_{t},s_{t+1})_{i}\}$ collected from various environments with arbitrary behaviour polices, our goal is to learn a single policy that maximizes the expected return $\mathcal{R}_{t}=\sum_{k\geq t}\gamma^{k-t}r_{k}$ on all considered environments and can be efficiently fine-tuned to new tasks, where $\gamma$ is the discount factor.

Considering the computational budget, we use 15 Atari games for pretraining and 5 games for out-of-distribution (OOD) experiments. The whole training process took around 12 days on A100 GPUs. Our offline dataset is derived from the the DQN-Replay dataset (Agarwal et al., ), which consists of $84\times 84$ grayscale images as observations and a full action space with 18 discrete actions.

## 4 Jointly-Optimized World-Action Model

In this section, we first detail the architecture of JOWA and the loss functions for joint optimization in section 4.1. Next, we introduce the provably efficient and parallelizable planning algorithm employed to compensate for the Q-value estimation error in section 4.2. Finally, we present the overall training and sample-efficient fine-tuning pipelines in section 4.3.

### 4.1 World-Action Model

#### 4.1.1 Architecture

Figure: Figure 1: 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.
Refer to caption: x1.png

Figure 1 illustrates JOWA’s architecture, which uses a transformer backbone (Vaswani, ) 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 are modeled via supervised learning using three heads: next observation token predictor $p_{o}$, reward predictor $p_{r}$, and termination predictor $p_{d}$. The Q-values head $h_{Q}$ learns the Q-function through TD-learning, based on implicit representations of historical trajectories. In the following sections, we refer to the {transformer, $p_{o}$, $p_{r}$, $p_{d}$} components as the ”world-part” with parameters $\theta$, and the {transformer, $h_{Q}$} components as the ”action-part” with parameters $\phi$.

We use VQ-VAE, a discrete autoencoder, to tokenize image observations, representing high-dimensional images as a sequence of $K$ tokens. The VQ-VAE is trained with an equally weighted combination of $L_{1}$ reconstruction loss, commitment loss (Van Den Oord et al., ), and perceptual loss (Esser et al., ). Then the transformer operates on the interleaved observation and action tokens represented as $(z_{0}^{1},\dots,z_{0}^{K},a_{0},\dots,z_{L}^{1},\dots,z_{L}^{K},a_{L})$, where $L$ 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 $t$, the world-part module models the following distributions:

$$ Dynamics predictor: $\displaystyle\hat{z}_{t}^{k}$ $\displaystyle\sim p_{o}\big{(}\hat{z}_{t}^{k}\mid f(z_{\leq t-1},z_{t}^{<k},a_ {\leq t-1},u)\big{)}$ (5) Reward predictor: $\displaystyle\hat{r}_{t}$ $\displaystyle\sim p_{r}\big{(}\hat{r}_{t}\mid f(z_{\leq t},a_{\leq t},u)\big{)}$ (6) Termination predictor: $\displaystyle\hat{d}_{t}$ $\displaystyle\sim p_{d}\big{(}\hat{d}_{t}\mid f(z_{\leq t},a_{\leq t},u)\big{)}$ (7) $$

where $f$ represents the transformer backbone, $u$ is the task ID to index the corresponding task embedding, $k\in\{1,\cdots,K\}$, and $t\in\{1,\cdots,L\}$.

To unify the category of loss functions for balanced training (Vandenhende et al., ), we convert scalar rewards to ternary quantities $\{-1,0,1\}$ using the sign function. This allows all three predictors to be optimized as classification problems by cross-entropy loss. Given $L$-timesteps segments sampled from the offline dataset, the loss function for the world-part module is formulated as:

$$ $\displaystyle\mathcal{L}_{\text{world}}(\theta)=\frac{1}{L}\sum_{t=1}^{L}\big{[}$ $\displaystyle\frac{1}{K}\sum_{k=1}^{K}-\ln p_{o}\big{(}z_{t}^{k}\mid f(z_{\leq t -1},z_{t}^{<k},a_{\leq t-1},u)\big{)}$ $\displaystyle-\ln p_{r}\big{(}r_{t}\mid f(z_{\leq t},a_{\leq t},u)\big{)}-\ln p _{d}\big{(}d_{t}\mid f(z_{\leq t},a_{\leq t},u)\big{)}\big{]}$ (8) $$

#### 4.1.3 Training of Action-part Module

We use CQL for offline TD-learning during pretraining. To ensure training stability and enhance scaling performance (Farebrother et al., ), we employ distributional TD-error (Bellemare et al., ) instead of the mean-square TD-error, maintaining consistency with the $\mathcal{L}_{\text{world}}$ loss category.

The action-part module computes the return distribution for all actions given an observation and historical information. For an observation-action pair $(o_{t},a_{t})$, the return distribution $\mathcal{Z}$ is formulated as: $\mathcal{Z}(o_{t},a_{t})=h_{Q}(f(z_{\leq t},a_{<t},u))[a_{t}]$. The value function $Q(o_{t},a_{t})$ is the mean of $\mathcal{Z}(o_{t},a_{t})$. Then the loss function for the action-part module is formulated as:

$$ $\mathcal{L}_{\text{action}}(\phi)=\alpha\left[\log\left(\sum_{a}\exp{\left(Q(o _{t},a)\right)}\right)-Q(o_{t},a_{t})\right]+\displaystyle D_{\mathrm{KL}} \left(\mathcal{T}^{*}\mathcal{Z}^{-}(o_{t},a_{t})\|\mathcal{Z}(o_{t},a_{t})\right)$ (9) $$

where $\mathcal{T}^{*}$ is the distributional Bellman optimality operator defined in Equation (1), and $\mathcal{Z}^{-}$ is the target distribution computed through a target Q-values head $h_{Q^{-}}$. We set $\alpha=0.1$ in experiments.

Therefore, the joint optimization objective for the world-action model is formulated as:

$$ $\mathcal{L}(\theta,\phi)=\beta\mathcal{L}_{\text{world}}(\theta)+\mathcal{L}_{ \text{action}}(\phi)$ (10) $$

with the coefficient $\beta>0$. We set $\beta=0.1$ in our experiments.

Both $\mathcal{L}_{\text{world}}$ and $\mathcal{L}_{\text{action}}$ back-propagate gradients to the transformer. Previous works observed that TD methods suffer from greater instability with larger model size (Kumar et al., ; Sokar et al., ). However, through jointly optimizing the world-action model, $\mathcal{L}_{\text{world}}$ serves as a regularizer to stabilize TD-learning in large-scale models. Moreover, the world-part module enables planning at decision time for optimal inference and sample-efficient transfer, detailed in the following sections.

### 4.2 Parallelizable planning at inference time

The world-part module enables planning at decision time to compensate for inaccurate Q-value estimates, allowing JOWA to consistently search out the optimal policy. We model this process as a tree search problem and present a practical, parallelizable search algorithm.

Given an estimated optimal Q-value $\hat{Q}^{*}$, a learned world model with dynamic predictor $\hat{P}$ and reward predictor $\hat{r}$, our objective is to find the optimal action $a^{*}_{0}$ maximizing the ground-truth optimal Q-value $Q^{*}$, starting from state $s_{0}$. To do so, we rewrite the Bellman optimality equation as:

$$ $Q^{*}(s_{0},a_{0})=\max_{\pi_{Q^{*}}}\operatorname*{\mathbb{E}}_{\begin{ subarray}{c}s_{1},\cdots,s_{H}\sim P\\ a_{1},\cdots,a_{H-1}\sim\pi_{Q^{*}}\end{subarray}}\left[\sum_{t=0}^{H-1}\gamma ^{t}r(s_{t},a_{t})+\gamma^{H}\max_{a_{H}}Q^{*}(s_{H},a_{H})\right]$ (11) $$

where $\pi_{Q^{*}}$ is the policy induced by the optimal Q-function. For Equation (11), the optimal policy is the greedy policy based on $Q^{*}$. The proof is provided in the Appendix A.

To derive the search objective function, we follow three steps: (i) replace the ground-truth functions in the right side of Equation (11) with estimated or learned functions. (ii) leverage the estimated optimal Q-function $\hat{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 $\Pi_{\hat{Q}^{*}}$, where $\forall\pi\in\Pi_{\hat{Q}^{*}},\forall s\in\mathcal{S},\forall a\notin\text{
top-}K(\hat{Q}^{*}(s,\cdot))$, we have $\pi(a|s)=0$. (iii) maximize over $a_{0}$ on both sides of Equation (11) to find the optimal initial action, considering the restriction in the second step. Finally, the resulting objective function is formulated as:

$$ $\max_{\pi\in\Pi_{\hat{Q}^{*}}}\operatorname*{\mathbb{E}}_{\begin{subarray}{c}s _{1},\cdots,s_{H}\sim\hat{P}\\ a_{0},\cdots,a_{H-1}\sim\pi\end{subarray}}\left[\underbrace{\sum_{t=0}^{H-1} \gamma^{t}\hat{r}(s_{t},a_{t})}_{\text{income-to-date}}+\underbrace{\gamma^{H} \max_{a_{H}}\hat{Q}^{*}(s_{H},a_{H})}_{\text{income-to-go}}\right]$ (12) $$

Detailed derivation is provided in the Appendix B. Then we show the error bound of search-based optimal Q-function in Theorem 14, with the formal theorem and its proof in Appendix C.

###### Theorem (Informal) 4.1 .

Denote the search-based optimal Q-function as $f(s,a)$. Assume the learned reward function $\hat{r}$ to be $L_{r}$-Lipschitz and the estimated optimal Q-function $\hat{Q}^{*}$ to be $L_{Q}$-Lipschitz. Assume the estimation errors of the learned state transition, reward, and Q-value are bounded by $\epsilon_{s},\epsilon_{r},\epsilon_{Q}$ respectively. Then we have the error between search-based optimal Q-value $f(s,a)$ and ground-truth optimal Q-value $Q^{*}(s,a)$ bounded as:

$$ $\left\|f(s,a)-Q^{*}(s,a)\right\|\leq\frac{1-\gamma^{H}}{1-\gamma}\epsilon_{r}+ \left(\frac{\gamma-\gamma^{H}}{1-\gamma}\epsilon_{r}+\gamma^{H}\epsilon_{Q} \right)\epsilon_{s}+\gamma^{H}\epsilon_{Q}$ (13) $$

Under the following condition, the search-based optimal Q-function $f$ has an upper error bound no greater than the estimated optimal Q-function $\hat{Q}^{*}$:

$$ $\frac{1}{1-\gamma}\epsilon_{r}+\left(\frac{\gamma}{1-\gamma}\frac{L_{r}-\gamma ^{H}L_{Q}}{1-\gamma^{H}}-\frac{L_{r}-L_{Q}}{1-\gamma}\frac{\gamma^{H}}{1- \gamma^{H}}\right)\epsilon_{s}\leq\epsilon_{Q}$ (14) $$

To optimize objective function (12), we interact with the imagined MDP induced by the learned world model using the constrained policy $\pi\in\Pi_{\hat{Q}^{*}}$ for $H$ steps, starting from $s_{0}$. We then compute the total income (income-to-date plus income-to-go) for all leaf nodes to identify the optimal path. The first edge of this path is the optimal initial action.

We implement this search using beam search, an efficient decoding algorithm common in NLP. At each timestep, we retain only $K$ states with the top-$K$ total income values. The horizon $H$ and beam width $K$ are hyper-parameters, with $K=1$ or $H=0$ degenerating to a $\hat{Q}^{*}$-based greedy policy. The algorithm calls the world model $K^{2}(H-1)+K$ times but only takes $H$ times as long as the forward propagation of the world model due to parallelizability across $K$ beams.

###### Theorem (Informal) 4.1 .

### 4.3 Training and Fine-tuning pipelines

Our multi-task offline RL pretraining consists of two stages:

- •
Stage 1: Sample trajectory segments from datasets. Train the VQ-VAE tokenizer using image observations. Train the world-part module using segments with loss (8) for $M_{1}$ steps.
- •
Stage 2: Freeze the VQ-VAE tokenizer. Sample segments and jointly optimize the world-action model with loss function (10) for $M_{2}$ steps.

We employ this two-stage training approach to stabilize and accelerate the overall training process. In our pretraining experiments, we set $M_{1}=250$k and $M_{2}=1.5$M, totaling $1.75$M gradient steps.

For sample-efficient transfer to unseen games, we adopt a similar two-stage offline fine-tuning pipeline, but unfreeze all components of JOWA in both two stages. In stage 1, we fine-tune JOWA using real data for 3k steps. In stage 2, we enable planning to synthesize high-quality data, using 3/4 batch of real and 1/4 batch of synthetic data with COMBO (4) as the action-part loss for 7k steps. Detailed fine-tuning protocol is shown in Appendix E.2.

## 5 Experiments

We design our experiments to answer the following questions: (1) How does JOWA perform on multi-games in low-data regime? (2) Can JOWA effectively leverage higher model capacity? (3) Does the pre-trained JOWA sample-efficiently transfer to new games?

### 5.1 Experimental Setup

##### Dataset.

We use the Atari trajectory dataset from Agarwal et al. (), which contains 50M transitions from each of 5 separate training runs. Due to the prohibitive long training time on full games, we select a subset of 20 games, maintaining the difficulty distribution of full games defined by Gulcehre et al. (). These 20 games are introduced in Appendix D. 15 of those games are used for training, and 5 games are held out for OOD generalization experiments. Following Lee et al. (); Kumar et al. (); Wu et al. (), we use data from 2 out of 5 training runs. To investigate performance in low-data regime, we uniformly draw 10% of transitions at random, as per Agarwal et al. (); Kumar et al. (), resulting in 10M transitions per game.

##### Training and Fine-tuning.

We implement our world-action model based on GPT-2 (Brown et al., ). 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 on A100 GPUs for 1.75M steps using the AdamW optimizer (Loshchilov & Hutter, ) with learning rates of $10^{-4}$ and $5\cdot 10^{-5}$ for two stages respectively, and batch size of 512 for both stages. For fine-tuning, we train for 10k gradient steps with a batch size of 32, learning rate of $5\cdot 10^{-5}$, and 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 beam 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., ), i.e. $(\textrm{score}-\textrm{score}_{\textrm{random}})/(\textrm{score}_{\textrm{
human}}-\textrm{score}_{\textrm{random}})$.
To create an aggregate comparison metric across all games, we use inter-quartile mean (IQM) of human-normalized scores, following evaluation best practices proposed in (Agarwal et al., ).

More details on hyperparameters, network architecture, algorithm implementation, and protocols for fine-tuning, and evaluation are provided in the Appendix E.

### 5.2 Baseline Methods

We compare JOWA with the following baselines: (i) Multi-Task BC (MTBC). Our method can be reduced to a transformer-based multi-task behavioral cloning (Pomerleau, ) agent by removing $\mathcal{L}_{\text{world}}$ and replacing the Q-values head with a policy head trained as a multi-class classification problem. We train MTBC models at 3 scales: 34M, 65M, 120M parameters. (ii) Multi-Game DT (MGDT). Lee et al. () used $\{o_{t+i},\hat{R}_{t+i},a_{t+i},r_{t+i}\}_{i=0}^{3}$ as input sequences, where $\hat{R}$ is the target return, and expert-conditioned return distribution to induce expert-level actions during evaluation. We train MGDT models at 2 scales: 40M and 200M parameters. (iii) Elastic DT (EDT). Based on MGDT, Wu et al. () removed the reward in the input sequences and dynamically selected history length to address challenges of trajectory stitching. We train EDT based on the architecture configuration of MGDT-200M. (iv) Scaled-QL (SQL). Kumar et al. () 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 ResNet-101 network. For fair comparison, all 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.1.

### 5.3 How does JOWA perform on multi-games in low-data regime?

**Table 2: 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.**
| Game | Random | Human | MTBC | MGDT | EDT | SQL | JOWA |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| 120M | 200M | 200M | 80M | 40M | 70M | 150M |  |  |  |
| Assault | 222.4 | 742.0 | 1203.5 | 1741.5 | 1915.5 | 2292 | 1578.3 | 1733.9 | 2302 |
| Atlantis | 12850 | 29028.1 | 37812.5 | 2565750 | 2108166.7 | 49100 | 41662.5 | 570862.5 | 2690387.5 |
| BeamRider | 363.9 | 16926.5 | 786 | 6011.3 | 4149.8 | 7023.5 | 880.3 | 2547.4 | 3498 |
| Berzerk | 123.7 | 2630.4 | 655 | 444.5 | 350 | 312.5 | 395 | 441.9 | 739 |
| Carnival | 0 | 3800 | 5560 | 2610 | 3678.8 | 1940 | 4685 | 4070 | 5316 |
| Centipede | 2090.9 | 12017 | 4046.6 | 4604 | 3389.9 | 3650.3 | 5669.3 | 4475.6 | 4677 |
| ChopperCommand | 811.0 | 7387.8 | 256.3 | 3300.8 | 3443.8 | 853.8 | 3118.8 | 2568.8 | 3812.5 |
| DemonAttack | 152.1 | 1971 | 2611.3 | 6549.4 | 3455.7 | 7936.5 | 1233.8 | 4584.4 | 3547.8 |
| NameThisGame | 2292.3 | 8049 | 6045 | 6610.5 | 7060 | 6461.7 | 7335.6 | 12706.9 | 11421 |
| Phoenix | 761.4 | 7242.6 | 1446.9 | 5120.5 | 5320 | 4961.7 | 1608.8 | 5065 | 5348 |
| Seaquest | 68.4 | 42054.7 | 120 | 2720 | 3160.4 | 650 | 1033.1 | 1490 | 2725 |
| SpaceInvaders | 148 | 1668.7 | 605 | 742.5 | 513.8 | 714.4 | 694.7 | 969.1 | 744.7 |
| StarGunner | 664.0 | 10250 | 1493.8 | 8625.5 | 9550 | 4728.6 | 5737.5 | 21231.3 | 18150 |
| TimePilot | 3568 | 5229.2 | 762.5 | 3866.7 | 2812.5 | 4072.7 | 3268.8 | 3831.3 | 3669 |
| Zaxxon | 32.5 | 9173.3 | 0 | 462.5 | 325.5 | 0 | 0 | 225 | 2163 |
| #Superhuman | 0 | N/A | 4 | 3 | 4 | 3 | 3 | 4 | 6 |
| Median HNS | 0.000 | 1.000 | 0.197 | 0.391 | 0.400 | 0.387 | 0.360 | 0.390 | 0.456 |
| IQM HNS | 0.000 | 1.000 | 0.329 | 0.498 | 0.502 | 0.420 | 0.371 | 0.476 | 0.789 |

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 and JOWA represent 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 augmentation during pretraining. JOWA-70M achieves an IQM human-normalized score of 47.6%, outperforming Scaled-QL-80M (42.0%). JOWA maintains its performance advantage when comparing median human-normalized scores. These results demonstrate JOWA’s sample efficiency in learning from heterogeneous offline data.

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?

Figure: Figure 2: Scaling trends on training set games.
Refer to caption: x2.png

Scaling law depicts the positive correlation between model capacity and performance. Following Lee et al. (); Kumar et al. (), we investigate JOWA’s ability to leverage higher capacity architectures. For comparison of scaling trends, we establish 2 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 2 shows the scaling trends for the 3 algorithms, which demonstrates that JOWA’s performance reliably increases as the model size grows. Although previous works observed that TD-learning suffer from greater instability with larger model size (Lee et al., ; Sokar et al., ), JOWA scales TD-learning through a scalable transformer architecture and an auxiliary regularization loss $\mathcal{L}_{\text{world}}$ 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.

**Table 3: Fine-tuning performance on unseen games using 5k transitions, measured in terms of DQN-normalized score, following Lee et al. (). See Table 17 in Appendix F for raw scores.**
| Game | MTBC | MGDT | EDT | SQL | JOWA | JOWA-150M |
| --- | --- | --- | --- | --- | --- | --- |
| 120M | 200M | 200M | 80M | 150M | (scratch) |  |
| Mean | 0.164 | 0.422 | 0.430 | 0.360 | 0.647 | 0.196 |
| Median | 0.215 | 0.354 | 0.325 | 0.284 | 0.615 | 0.173 |
| IQM | 0.205 | 0.377 | 0.380 | 0.355 | 0.647 | 0.181 |

### 5.5 Can JOWA sample-efficiently transfer to new games?

Rapid adaptation to downstream tasks is a natural and well-motivated benefit of pretraining. In this section, we study how the pretrained world-action model enables rapid and sample-efficient fine-tuning on new games. To investigate this question, we fine-tune pretrained agents on 5 held-out games using uniformly subsampled 5k expert-level transitions (from last 20% of DQN-Replay (Agarwal et al., )) per game as the benchmark. These tiny amounts of transitions corresponding to approximately 4 trajectories from expert-level DQN-Replay per fine-tuned game on average, which is similar to the settings of few-shot learning and extremely challenging. For comparison, we fine-tune the largest agents of baselines alongside a no-pretrained baseline, JOWA-150M (scratch) trained on each held-out game from scratch. All pretrained agents are fine-tuned for 10k steps while JOWA-150M (scratch) is trained for 100k gradient steps. Detailed fine-tuning protocol is shown in Appendix E.2. We report results in terms of DQN-normalized score, following Lee et al. (); Kumar et al. ().

We summarize the fine-tuning results in Table 3. This table shows that the fine-tuned JOWA-150M attains 64.7% IQM DQN-normalized score across 5 held-out games, outperforming baselines by 34.7% on averange. These results underscore JOWA’s capacity for rapid and sample-efficient transfer to novel games, highlighting the efficacy of its learned general-purpose representation and decision-making capabilities.

### 5.6 Ablation Study

In this section, we conduct a series of controlled ablation studies to evaluate the impact of key design choices in JOWA, including planning, different training losses, and architecture designs. These experiments aim to provide empirical evidence to support our designs and offer valuable insights for future research in this domain.

**Table 4: Performance on the 15 Atari games for 3 variants of JOWA evaluated with or without planning reported in terms of the IQM human-normalized score.**
|  | JOWA-150M | JOWA-70M | JOWA-40M |
| --- | --- | --- | --- |
| without planning | 50.8% | 45.1% | 27.1% |
| with planning | 78.9% | 47.6% | 37.1% |
| improvement ($\uparrow$) | +28.1% | +2.5% | +10.0% |

##### Effects of planning at decision time.

We evaluate our 3 variants models with or without planning on the pretrained games and report the IQM aggregated human-normalized score across 15 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 demonstrates scaling trends with the growth of parameter size, but at a slower scaling rate.

The subsequent ablation experiments require training from scratch in multi-game regime. For time-saving, we consider a subset of 6 games in the following experiments: Assault, Carnival, Centipede, NameThisGame, Phoenix, SpaceInvaders. We train all models with 10% randomly subsampled data per game for 1M gradient steps and fix the parameter size to 40M.

**Table 5: The mean, median, and IQM human-normalized score on the 6 Atari games for various training choices. See Tabel 18 in Appendix F for raw scores.**
| Game | Origin | Different training losses | Q-heads ensemble | No task | Synthetic data |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| No CQL | No $\mathcal{L}_{\text{world}}$ | $\operatorname{sg}(\mathcal{L}_{\text{action}})$ | MSE | Equal | Random | embedding | in pretraining |  |  |
| 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 |

##### Different training losses.

We train models with each of the following 5 loss functions to investigate the impact of each loss term: (i) original loss defined in Equation (10), (ii) No CQL regularization in $\mathcal{L}_{\text{action}}$, which means no conservative constraints, (iii) No $\mathcal{L}_{\text{world}}$, which means training a model-free transformer-based critic network using $\mathcal{L}_{\text{action}}$, (iv) gradients of $\mathcal{L}_{\text{action}}$ not back-propagated to transformer backbone, denoted as $\operatorname{sg}(\mathcal{L}_{\text{action}})$ for short, which means optimizing the world model and critic network separately, (v) mean square error (MSE) instead of distributional TD-loss. The results of different training losses are shown in the third main column of Table 5.

Observe that the $\operatorname{sg}(\mathcal{L}_{\text{action}})$ and MSE loss fail to train qualified multi-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 $\alpha$ (tested with $\alpha\in\{0.01,0.05,0.1\}$), resulting in extreme Q-values for in-domain state-actions and low Q-values for OOD actions. The fail of $\operatorname{sg}(\mathcal{L}_{\text{action}})$ underscores the importance of joint-optimization in the world-action model. Comparing the original loss with the no $\mathcal{L}_{\text{world}}$ 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 a large margin, indicating the effectiveness of every loss terms.

##### Ensemble of Q-value heads.

We run JOWA-40M with ensembled Q-heads, where the return distribution is a 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 3 and compare 2 ensemble approaches: equal weights or random weights like REM (Agarwal et al., ). We report the results in the forth main column of Table 5. Although the equal weighted ensemble slightly outperforms random weights, we have not observed significant improvements of multi Q-heads over single Q-head. We leave the better way of distributional-Q 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 fifth main column of Table 5. Observe that the addition of task embedding improves 21.9% and 35.7% for mean and median human-normalized score respectively.

##### Effects of synthetic data in pretraining.

By default, we do not employ planning or synthetic data during pretraining due to the significant increase in training time (approximately 10 times slower, extending training to months). For this ablation study, we enable synthetic data in pretraining for JOWA-40M and train on 6 games for 1M steps with the same batch size and all other hyperparameters, which takes 31 days. Specifically, we sample batch of 4-frames segments and interact in the imagined MDP with $\epsilon$-greedy policy to synthesis the last 4 steps. Then we optimize the world-action model with half batch of real data and half batch of synthetic data using COMBO (Yu et al., ) loss in Equation (4) as $\mathcal{L}_{\text{action}}$. The results are shown in the last main column of Table 5.

Suprisingly, we observe negative gains from incorporating synthetic data. We hypothesize twofold reasons: (i) accumulation of inference errors over steps causes later synthetic steps to deviate significantly from the ground-truth distribution, (ii) COMBO’s over-penalization 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 for utilizing synthetic data during multi-game RL pretraining is left for future work.

## 6 Conclusion

In this work, we introduce JOWA: Jointly-Optimized 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 RL methods in multi-games regime. Furthermore, by training a large-capacity model on a 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., ), as this would take two 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.

## Appendix A Proof of Equation ( 11 )

Given the initial state-action pair $(s_{0},a_{0})$, the bellman expectation equation (Sutton, ) is written as:

$$ $\displaystyle Q^{\pi}(s_{0},a_{0})$ $\displaystyle=\operatorname*{\mathbb{E}}_{s_{1}\sim P(\cdot|s_{0},a_{0})}\left [r(s_{0},a_{0})+\gamma V^{\pi}(s_{1})\right]$ $\displaystyle=\operatorname*{\mathbb{E}}_{\begin{subarray}{c}s_{1},\cdots,s_{H }\sim P\\ a_{1},\cdots,a_{H-1}\sim\pi\end{subarray}}\left[\sum_{t=0}^{H-1}\gamma^{t}r(s_ {t},a_{t})+\gamma^{H}V^{\pi}(s_{H})\right]$ (15) $$

According to the definition of optimal Q-value: $Q^{*}(s_{0},a_{0})=\max_{\pi}Q^{\pi}(s_{0},a_{0})$, we have:

$$ $\displaystyle Q^{*}(s_{0},a_{0})$ $\displaystyle=\max_{\pi}\operatorname*{\mathbb{E}}_{\begin{subarray}{c}s_{1}, \cdots,s_{H}\sim P\\ a_{1},\cdots,a_{H-1}\sim\pi\end{subarray}}\left[\sum_{t=0}^{H-1}\gamma^{t}r(s_ {t},a_{t})+\gamma^{H}V^{\pi}(s_{H})\right]$ $\displaystyle=\max_{\pi}\operatorname*{\mathbb{E}}_{\begin{subarray}{c}s_{1}, \cdots,s_{H}\sim P\\ a_{1},\cdots,a_{H-1}\sim\pi\end{subarray}}\left[\sum_{t=0}^{H-1}\gamma^{t}r(s_ {t},a_{t})+\gamma^{H}\max_{\pi}V^{\pi}(s_{H})\right]$ $\displaystyle=\max_{\pi}\operatorname*{\mathbb{E}}_{\begin{subarray}{c}s_{1}, \cdots,s_{H}\sim P\\ a_{1},\cdots,a_{H-1}\sim\pi\end{subarray}}\left[\sum_{t=0}^{H-1}\gamma^{t}r(s_ {t},a_{t})+\gamma^{H}V^{*}(s_{H})\right]$ $\displaystyle=\max_{\pi}\operatorname*{\mathbb{E}}_{\begin{subarray}{c}s_{1}, \cdots,s_{H}\sim P\\ a_{1},\cdots,a_{H-1}\sim\pi\end{subarray}}\left[\sum_{t=0}^{H-1}\gamma^{t}r(s_ {t},a_{t})+\gamma^{H}\max_{a_{H}}Q^{*}(s_{H},a_{H})\right]$ (16) $$

For Equation (16), we derive the optimal policy $\pi$ is the greedy policy selecting actions with the greatest Q*-value as follows:

$$ $\displaystyle Q^{*}(s_{0},a_{0})$ $\displaystyle=\max_{\pi}\operatorname*{\mathbb{E}}_{\begin{subarray}{c}s_{1}, \cdots,s_{H}\sim P\\ a_{1},\cdots,a_{H-1}\sim\pi\end{subarray}}\left[\sum_{t=0}^{H-1}\gamma^{t}r(s_ {t},a_{t})+\gamma^{H}\max_{a_{H}}Q^{*}(s_{H},a_{H})\right]$ $\displaystyle=\max_{\pi}\operatorname*{\mathbb{E}}_{\begin{subarray}{c}s_{1}, \cdots,s_{H-1}\sim P\\ a_{1},\cdots,a_{H-2}\sim\pi\end{subarray}}\left[\sum_{t=0}^{H-2}\gamma^{t}r(s_ {t},a_{t})+\gamma^{H-1}\max_{a_{H-1}}\operatorname*{\mathbb{E}}_{s_{H}\sim P} \left[r(s_{H-1},a_{H-1})+\gamma\max_{a_{H}}Q^{*}(s_{H},a_{H})\right]\right]$ $\displaystyle=\max_{\pi}\operatorname*{\mathbb{E}}_{\begin{subarray}{c}s_{1}, \cdots,s_{H-1}\sim P\\ a_{1},\cdots,a_{H-2}\sim\pi\end{subarray}}\left[\sum_{t=0}^{H-2}\gamma^{t}r(s_ {t},a_{t})+\gamma^{H-1}\max_{a_{H-1}}Q^{*}(s_{H-1},a_{H-1})\right]$ (17) $$

Therefore, we have $a_{H-1}=\arg\max_{a}Q^{*}(s_{H-1},a)$. Similarly, continuing to use dynamic programming on Equation (17), we finally get: $a_{i}=\arg\max_{a}Q^{*}(s_{i},a),(i=1,2,\cdots,H-1)$. Thus we claim that the optimal policy is induced by the optimal value and use the symbol $\pi_{Q^{*}}$ instead of $\pi$ in Equation (16) to obtain Equation (11).

## Appendix B Derivation from Equation ( 11 ) to Equation ( 12 )

Given the learned dynamic preditor $\hat{P}$, reward predictor $\hat{r}$, and estimated optimal Q-value $\hat{Q}^{*}$, we first substitute these three functions for the ground-truth functions in the right side of Equation (11):

$$ $\max_{\pi_{\hat{Q}^{*}}}\operatorname*{\mathbb{E}}_{\begin{subarray}{c}s_{1}, \cdots,s_{H}\sim\hat{P}\\ a_{1},\cdots,a_{H-1}\sim\pi_{\hat{Q}^{*}}\end{subarray}}\left[\sum_{t=0}^{H-1} \gamma^{t}\hat{r}(s_{t},a_{t})+\gamma^{H}\max_{a_{H}}\hat{Q}^{*}(s_{H},a_{H})\right]$ (18) $$

However, due to the estimation error between learned functions and ground-truth functions, $\hat{Q}^{*}(s_{t},a_{t})$ is typically not equal to $\operatorname*{\mathbb{E}}_{s_{t+1}\sim\hat{P}}\left[\hat{r}(s_{t},a_{t})+
\gamma\max_{a}\hat{Q}^{*}(s_{t+1},a)\right]$. Therefore, instead of using dynamic programming to derive that the optimal policy is the greedy policy as in Proof A, we have to use search to solve formula (18).

However, we can use the conclusion derived with ground-truth functions to make assumptions to reduce the policy space for search. We assume that the optimal policy for formula (18) maps states to actions with top-K highest Q-values. Denote the constrained policy space as $\Pi_{\hat{Q}^{*}}$, where $\forall\pi\in\Pi_{\hat{Q}^{*}},\forall s\in\mathcal{S},\forall a\notin\text{
top-}K(\hat{Q}^{*}(s,\cdot))$, we have $\pi(a|s)=0$. we make the following assumption: the optimal policy of formula (18) is in the constrained policy space, *i.e.,*$\pi^{*}_{\hat{Q}^{*}}\in\Pi_{\hat{Q}^{*}}$. Under this assumption, formula (18) is equivalent to:

$$ $\max_{\pi\in\Pi_{\hat{Q}^{*}}}\operatorname*{\mathbb{E}}_{\begin{subarray}{c}s _{1},\cdots,s_{H}\sim\hat{P}\\ a_{1},\cdots,a_{H-1}\sim\pi\end{subarray}}\left[\sum_{t=0}^{H-1}\gamma^{t}\hat {r}(s_{t},a_{t})+\gamma^{H}\max_{a_{H}}\hat{Q}^{*}(s_{H},a_{H})\right]$ (19) $$

Then we maximize formula (19) over $a_{0}$ to find the optimal initial action. Considering the above assumption, $a_{0}\in\text{top-}K(\hat{Q}^{*}(s_{0},\cdot))$. Thus we have the objective for search as:

$$ $\max_{\pi\in\Pi_{\hat{Q}^{*}}}\operatorname*{\mathbb{E}}_{\begin{subarray}{c}s _{1},\cdots,s_{H}\sim\hat{P}\\ a_{0},\cdots,a_{H-1}\sim\pi\end{subarray}}\left[\sum_{t=0}^{H-1}\gamma^{t}\hat {r}(s_{t},a_{t})+\gamma^{H}\max_{a_{H}}\hat{Q}^{*}(s_{H},a_{H})\right]$ (20) $$

## Appendix C Upper bound of search-based Q-value estimation

Let function $f(s_{0},a_{0})$ be equal to formula (19) and let $\pi^{*}_{\hat{Q}^{*}}$ be the optimal policy of formula (19). We have:

$$ $f(s_{0},a_{0})=\operatorname*{\mathbb{E}}_{\begin{subarray}{c}s_{1},\cdots,s_{ H}\sim\hat{P}\\ a_{1},\cdots,a_{H-1}\sim\pi^{*}_{\hat{Q}^{*}}\end{subarray}}\left[\sum_{t=0}^{ H-1}\gamma^{t}\hat{r}(s_{t},a_{t})+\gamma^{H}\max_{a_{H}}\hat{Q}^{*}(s_{H},a_{ H})\right]$ (21) $$

We make the following assumption similar to EfficientZero-v2 (Wang et al., ):

###### Assumption C.1 .

Assume the state transition, reward, and Q-value estimations error are upper bounded by $\epsilon_{s},\epsilon_{r},\epsilon_{Q}$ respectively. The error bound of each estimation is formulated as:

$$ $\displaystyle\max_{n\in[N],t\in[H(n)]}\mathbb{E}\left[\|\hat{s}_{t}-s_{t}\| \right]\leq\epsilon_{s}$ (22) $\displaystyle\max_{n\in[N],t\in[H(n)]}\mathbb{E}\left[\|\hat{r}(s_{t})-r(s_{t} )\|\right]\leq\epsilon_{r}$ (23) $\displaystyle\max_{n\in[N],t\in[H(n)]}\mathbb{E}\left[\|\hat{Q}^{*}(s_{t})-Q^{ *}(s_{t})\|\right]\leq\epsilon_{Q}$ (24) $$

###### Theorem C.2 .

Define $s_{t},a_{t}$ to be the states and actions resulting from current policy using ground-truth dynamics $P$ and reward function $r$ and similarly define $s_{t}^{{}^{\prime}},a_{t}^{{}^{\prime}}$ using learned functions $\hat{P}$ and $\hat{r}$. Assume the learned reward function $\hat{r}$ to be $L_{r}$-Lipschitz and the estimated optimal Q-function $\hat{Q}^{*}$ to be $L_{Q}$-Lipschitz. Assume the estimation errors of learned functions are bounded as in Assumption C.1. Then we have the error between search-based Q-value estimation $f(s_{0},a_{0})$ and ground-truth Q-value $Q^{*}(s_{0},a_{0})$ bounded as:

$$ $\|f(s_{0},a_{0})-Q^{*}(s_{0},a_{0})\|\leq\frac{1-\gamma^{H}}{1-\gamma}\epsilon _{r}+\left(\frac{\gamma-\gamma^{H}}{1-\gamma}\epsilon_{r}+\gamma^{H}\epsilon_{ Q}\right)\epsilon_{s}+\gamma^{H}\epsilon_{Q}$ (25) $$

###### Proof.

$$ $\displaystyle\|f(s_{0},a_{0})-Q^{*}(s_{0},a_{0})\|$ $\displaystyle=$ $\displaystyle\left\|\operatorname*{\mathbb{E}}\left[\sum_{t=0}^{H-1}\gamma^{t} \hat{r}(s_{t}^{{}^{\prime}},a_{t}^{{}^{\prime}})+\gamma^{H}\max_{a_{H}}\hat{Q} ^{*}(s_{H}^{{}^{\prime}},a_{H})\right]-\operatorname*{\mathbb{E}}\left[\sum_{t =0}^{H-1}\gamma^{t}r(s_{t},a_{t})+\gamma^{H}\max_{a_{H}}Q^{*}(s_{H},a_{H}) \right]\right\|$ $\displaystyle\leq$ $\displaystyle\operatorname*{\mathbb{E}}\left[\left\|\hat{r}(s_{0},a_{0})-r(s_{ 0},a_{0})+\sum_{t=1}^{H-1}\gamma^{t}\left(\hat{r}(s_{t}^{{}^{\prime}},a_{t}^{{ }^{\prime}})-r(s_{t},a_{t})\right)+\gamma^{H}\left(\max_{a_{H}}\hat{Q}^{*}(s_{ H}^{{}^{\prime}},a_{H})-\max_{a_{H}}Q^{*}(s_{H},a_{H})\right)\right\|\right]$ $\displaystyle\leq$ $\displaystyle\operatorname*{\mathbb{E}}\left[\left\|\hat{r}(s_{0},a_{0})-r(s_{ 0},a_{0})\right\|\right]+\sum_{t=1}^{H-1}\gamma^{t}\left\|\hat{r}(s_{t}^{{}^{ \prime}},a_{t}^{{}^{\prime}})-r(s_{t},a_{t})\right\|+\gamma^{H}\operatorname*{ \mathbb{E}}\left[\left\|\max_{a_{H}}\hat{Q}^{*}(s_{H}^{{}^{\prime}},a_{H})- \max_{a_{H}}Q^{*}(s_{H},a_{H})\right\|\right]$ $\displaystyle\leq$ $\displaystyle\epsilon_{r}+\sum_{t=1}^{H-1}\gamma^{t}\left\|\hat{r}(s_{t}^{{}^{ \prime}},a_{t}^{{}^{\prime}})-r(s_{t},a_{t})\right\|+\gamma^{H}\operatorname*{ \mathbb{E}}\left[\left\|\max_{a_{H}}\hat{Q}^{*}(s_{H}^{{}^{\prime}},a_{H})- \max_{a_{H}}Q^{*}(s_{H},a_{H})\right\|\right]$ (26) $$

For the second term in inequality (26):

$$ $\displaystyle\left\|\hat{r}(s_{t}^{{}^{\prime}},a_{t}^{{}^{\prime}})-r(s_{t},a _{t})\right\|$ $\displaystyle=$ $\displaystyle\left\|\hat{r}(s_{t}^{{}^{\prime}},a_{t}^{{}^{\prime}})-\hat{r}(s _{t},a_{t})+\hat{r}(s_{t},a_{t})-r(s_{t},a_{t})\right\|$ $\displaystyle\leq$ $\displaystyle\left\|\hat{r}(s_{t}^{{}^{\prime}},a_{t}^{{}^{\prime}})-\hat{r}(s _{t},a_{t})\right\|+\left\|\hat{r}(s_{t},a_{t})-r(s_{t},a_{t})\right\|$ $\displaystyle\leq$ $\displaystyle L_{r}\left\|s_{t}^{{}^{\prime}}-s_{t}\right\|+\epsilon_{r}$ $\displaystyle\leq$ $\displaystyle L_{r}\epsilon_{s}+\epsilon_{r}$ (27) $$

For the third term in inequality (26), let $a_{H}^{1}=\arg\max_{a_{H}}\hat{Q}^{*}(s_{H}^{{}^{\prime}},a_{H})$ and $a_{H}^{2}=\arg\max_{a_{H}}Q^{*}(s_{H},a_{H})$. Then we have:

$$ $\displaystyle\operatorname*{\mathbb{E}}\left[\left\|\max_{a_{H}}\hat{Q}^{*}(s_ {H}^{{}^{\prime}},a_{H})-\max_{a_{H}}Q^{*}(s_{H},a_{H})\right\|\right]$ $\displaystyle=$ $\displaystyle\operatorname*{\mathbb{E}}\left[\left\|\hat{Q}^{*}(s_{H}^{{}^{ \prime}},a_{H}^{1})-Q^{*}(s_{H},a_{H}^{2})\right\|\right]$ $\displaystyle=$ $\displaystyle\operatorname*{\mathbb{E}}\left[\left\|\hat{Q}^{*}(s_{H}^{{}^{ \prime}},a_{H}^{1})-\hat{Q}^{*}(s_{H},a_{H}^{2})+\hat{Q}^{*}(s_{H},a_{H}^{2})- Q^{*}(s_{H},a_{H}^{2})\right\|\right]$ $\displaystyle\leq$ $\displaystyle\operatorname*{\mathbb{E}}\left[\left\|\hat{Q}^{*}(s_{H}^{{}^{ \prime}},a_{H}^{1})-\hat{Q}^{*}(s_{H},a_{H}^{2})\right\|\right]+\operatorname* {\mathbb{E}}\left[\left\|\hat{Q}^{*}(s_{H},a_{H}^{2})-Q^{*}(s_{H},a_{H}^{2}) \right\|\right]$ $\displaystyle\leq$ $\displaystyle L_{Q}\left\|s_{H}^{{}^{\prime}}-s_{H}\right\|+\epsilon_{Q}$ $\displaystyle\leq$ $\displaystyle L_{Q}\epsilon_{s}+\epsilon_{Q}$ (28) $$

Substitute inequalities (27) and (28) into (26):

$$ $\displaystyle\|f(s_{0},a_{0})-Q^{*}(s_{0},a_{0})\|$ $\displaystyle\leq$ $\displaystyle\epsilon_{r}+\sum_{t=1}^{H-1}\gamma^{t}(L_{r}\epsilon_{s}+ \epsilon_{r})+\gamma^{H}(L_{Q}\epsilon_{s}+\epsilon_{Q})$ $\displaystyle=$ $\displaystyle\frac{1-\gamma^{H}}{1-\gamma}\epsilon_{r}+\left(\frac{\gamma- \gamma^{H}}{1-\gamma}\epsilon_{r}+\gamma^{H}\epsilon_{Q}\right)\epsilon_{s}+ \gamma^{H}\epsilon_{Q}$ (29) $$

∎

###### Assumption C.1 .

###### Theorem C.2 .

###### Proof.

##### Analysis.

We expect the search-based Q-values $f(s,a)$ have an upper error bound no greater than the estimated Q-values $\hat{Q}^{*}(s,a)$, which is formulated as the following inequality:

$$ $\frac{1-\gamma^{H}}{1-\gamma}\epsilon_{r}+\left(\frac{\gamma-\gamma^{H}}{1- \gamma}\epsilon_{r}+\gamma^{H}\epsilon_{Q}\right)\epsilon_{s}+\gamma^{H} \epsilon_{Q}\leq\epsilon_{Q}$ (30) $$

Based on inequality (30), we derive the following condition:

$$ $\frac{1}{1-\gamma}\epsilon_{r}+\left(\frac{\gamma}{1-\gamma}\frac{L_{r}-\gamma ^{H}L_{Q}}{1-\gamma^{H}}-\frac{L_{r}-L_{Q}}{1-\gamma}\frac{\gamma^{H}}{1- \gamma^{H}}\right)\epsilon_{s}\leq\epsilon_{Q}$ (31) $$

The inequality (31) means that if the weighted sum of rewards estimation error $\epsilon_{r}$ and state transition estimation error $\epsilon_{s}$ are less than or equal to the Q-values estimation error $\epsilon_{Q}$, then the search-based optimal Q-values have a lower upper-bound of error than estimated optimal Q-values.

## Appendix D Games

We select 20 Atari games maintaining the difficulty distribution of full Atari 2600 games defined by Gulcehre et al. (), 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 experiments. The 15 training games are: Phoenix, Centipede, SpaceInvaders, Carnival, NameThisGame, Assault, Atlantis, DemonAttack, BeamRider, ChopperCommand, 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 6.

**Table 6: Atari Games: Name, Game difficulty, Action Space, and Type.**
| Game | Difficulty | Action Space | Type |
| --- | --- | --- | --- |
| Assault | Medium | 7 | Train |
| Atlantis | Hard | 18 | Train |
| BeamRider | Medium | 9 | Train |
| Berzerk | Hard | 18 | Train |
| Carnival | Medium | 6 | Train |
| Centipede | Medium | 18 | Train |
| ChopperCommand | Easy | 18 | Train |
| DemonAttack | Easy | 6 | Train |
| Gravitar | Easy | 18 | Fine-tune |
| MsPacman | Medium | 9 | Fine-tune |
| NameThisGame | Easy | 6 | Train |
| Phoenix | Easy | 8 | Train |
| Pong | Medium | 6 | Fine-tune |
| Robotank | Medium | 18 | Fine-tune |
| Seaquest | Easy | 18 | Train |
| SpaceInvaders | Easy | 6 | Train |
| StarGunner | Medium | 18 | Train |
| TimePilot | Easy | 10 | Train |
| YarsRevenge | Medium | 18 | Fine-tune |
| Zaxxon | Easy | 18 | Train |

## Appendix E Experimental Details

### E.1 Implement Details

#### E.1.1 JOWA

We implement JOWA based on the codes of IRIS (Micheli et al., )(^1^11https://github.com/eloialonso/iris). We train the tokenizer VQ-VAE using the following loss function:

$$ $\mathcal{L}(E,D,\mathcal{E})=\|x-D(z)\|_{1}+\|\mathrm{sg}(E(x))-\mathcal{E}(z) \|^{2}_{2}+\|\mathrm{sg}(\mathcal{E}(z))-E(x)\|^{2}_{2}+\mathcal{L}_{ perceptual}(x,D(z))$ $$

where $E,D,\mathcal{E}$ are encoder, decoder, and embedding table respectively. $\mathrm{sg}(\cdot)$ is the stop-gradient operator. The last term is the perceptual loss (Johnson et al., ). We list the hyperparameters of VQ-VAE in Table 8 and 8. After the first stage of pretraining, the VQ-VAE is frozen.

**Table 7: Encoder / Decoder hyperparameters. We list the hyperparameters for the encoder, the same ones apply for the decoder.**
| Hyperparameter | Value |
| --- | --- |
| Frame dimensions (h, w) | $84\times 84$ |
| Layers | 3 |
| Residual blocks per layer | 2 |
| Channels in convolutions | 64 |
| Self-attention layers at resolution | 6 / 12 |

In addition to the vocabulary embedding and position embedding, we add a learnable task embedding for observation tokens and action tokens respectively. Our transformer are based on minGPT(^2^22https://github.com/karpathy/minGPT) with FlashAttention (Dao et al., ) for acceleration. The hyperparameters of JOWA’s transformer backbone are listed in Table 10 and 10.

**Table 9: Same hyperparameters of transformer for 3 JOWA variants.**
| Hyperparameter | Value |
| --- | --- |
| max sequence tokens | 296 |
| dropout rate | 0.1 |

**Table 11: Hyperparameters of Q-heads for 3 JOWA variants.**
| Model | Layers | MLP Hidden dimension | Number of heads | Dropout |
| --- | --- | --- | --- | --- |
| JOWA-40M | 3 | 768 | 1 | 0.01 |
| JOWA-70M | 3 | 1024 | 1 | 0.01 |
| JOWA-150M | 3 | 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. (). The hyperparameters of Q-heads for JOWA are shown in Table 11. The training hyperparameters of JOWA are shown in Table 12.

**Table 12: Training hyperparameters of JOWA.**
| Hyperparameter | Value |
| --- | --- |
| Optimizer (VQ-VAE) | Adam |
| Optimizer (except VQ-VAE) | AdamW |
| Learning rate (VQ-VAE) | 0.0001 |
| Learning rate (except VQ-VAE, stage 1) | 0.0001 |
| Learning rate (except VQ-VAE, stage 2) | 0.00005 |
| Batch size (VQ-VAE) | 2048 |
| Batch size (except VQ-VAE) | 512 |
| Weight decay (except VQ-VAE) | 0.01 |
| Gradient clip | 1.0 |
| Discount factor ($\gamma$) | 0.99 |
| Target Q update frequency | 1000 |
| Distributional Q | [-10, 30] |
| Number of atoms | 51 |
| Coefficient of CQL ($\alpha$) | 0.1 |
| Coefficient of $\mathcal{L}_{\text{world}}$ ($\beta$) | 0.1 |

**Table 13: Hyperparameters of transformer for 2 MGDT variants.**
| Model | Layers | Hidden size | Heads |
| --- | --- | --- | --- |
| MGDT-40M | 6 | 768 | 12 |
| MGDT-200M | 10 | 1280 | 20 |

**Table 14: Training hyperparameters of Scaled-QL.**
| Hyperparameter | Value |
| --- | --- |
| Optimizer | Adam |
| Learning rate | 0.0002 |
| Batch size | 512 |
| Gradient clip | 1.0 |
| Discount factor ($\gamma$) | 0.99 |
| Target Q update frequency | 2000 |
| Distributional Q | [-20, 20] |
| Number of atoms | 51 |
| Coefficient of CQL ($\alpha$) | 0.05 |
| n-step returns | 3 |

#### E.1.2 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 as a 18-class classification problem. All hyperparameters are kept the same as JOWA.

#### E.1.3 EDT

We use the official code for EDT(^3^33https://github.com/kristery/Elastic-DT). We implement EDT-200M based on the architecture configuration of MGDT-200M, which is shown in Table 13. We change the batch size to 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 $\{o_{t+i},R_{t+i},a_{t+i},r_{t+i}\}_{i=0}^{3}$ as the input sequences, remove the expectile regression loss $\mathcal{L}_{\max}$ and observation prediction loss $\mathcal{L}_{\text{observation}}$, and add the reward prediction loss to rewrite the codes of EDT into MGDT. We enable data augmentation (random cropping and random rotation) for MGDT. The hyperparameters of transformer for two MGDT variants are listed in Table 13.

#### E.1.5 Scaled-QL

We implement a pytorch version of Scaled-QL from scratch, referring to the jax version of its official preliminary codes(^4^44https://tinyurl.com/scaled-ql-code). We use ResNet-101 as 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 a 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 $\frac{e}{\|e\|^{2}}$. The training hyperparameters of Scaled-QL are listed in Table 15.

For fair comparison, all methods are 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., ) (last 20% of the original dataset) for each held-out game. Each game was fine-tuned separately to measure the model’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 10k gradient steps. For reporting results, we report the performance of the agent snapshot that obtain the highest score during fine-tuning.

For JOWA in the second fine-tuning stage, we set 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 2 steps. Then we update JOWA with 3/4 batch (24) of real data and 1/4 batch (8) of synthetic data using COMBO loss rather than CQL loss as the $\mathcal{L}_{\text{action}}$. For other baselines, we enable random cropping and random rotation for data augmentation.

### 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. (); Kumar et al. (). Following standard protocols on Atari, we evaluate a noised version of the policy with an epsilon-greedy scheme, with $\epsilon_{\text{eval}}=0.001$. The evaluation settings of Atari are shown in Table 15.

For the expert action inference of MGDT and EDT, we set the inverse temperature $\kappa$ to 10. For the planning of JOWA, we set the planning horizon $H$ to 2 for all games. The beam width are set according to the size of valid action space of each game. Specifically, we set beam width $K$ to 2 if the valid action space size is less than 10, otherwise we set $K$ in $\{3,4\}$. The planning hyperparameters for each game are shown in Table 16.

**Table 16: Planning hyperparameters of JOWA during evaluation.**
| Game | planning horizon | beam width | Action space |
| --- | --- | --- | --- |
| Assault | 2 | 2 | 7 |
| Atlantis | 2 | 3 | 18 |
| BeamRider | 2 | 2 | 9 |
| Berzerk | 2 | 4 | 18 |
| Carnival | 2 | 2 | 6 |
| Centipede | 2 | 4 | 18 |
| ChopperCommand | 2 | 4 | 18 |
| DemonAttack | 2 | 2 | 6 |
| NameThisGame | 2 | 2 | 6 |
| Phoenix | 2 | 2 | 8 |
| Seaquest | 2 | 3 | 18 |
| SpaceInvaders | 2 | 2 | 6 |
| StarGunner | 2 | 3 | 18 |
| TimePilot | 2 | 3 | 10 |
| Zaxxon | 2 | 3 | 18 |

## Appendix F Raw scores

We summarize the raw scores of fine-tuning experiments and ablation studies in Table 17 and 18 respectively.

**Table 17: Offline fine-tuning performance on unseen games using 5k transitions, measured in terms of DQN-normalized score, following Lee et al. (); Kumar et al. ().**
| Game | Random | DQN | MTBC | MGDT | EDT | SQL | JOWA | JOWA-150M |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| 120M | 200M | 200M | 80M | 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 18: Raw scores on the 6 Atari games for various training choices. The mean, median, and IQM human-normalized score are shown in the last 3 rows and the best scores are markded in bold.**
| Game | Origin | Different training losses | Q-heads ensemble | No task | Synthetic data |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| No CQL | No $\mathcal{L}_{\text{world}}$ | $\operatorname{sg}(\mathcal{L}_{\text{action}})$ | MSE | Equal | Random | embedding | in pretraining |  |  |
| Assault | 1423.5 | 857.9 | 1650 | 452.7 | 637.6 | 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 | 5568.7 | 4097.3 | 5592.5 | 6450 | 3725.4 | 3253 |
| NameThisGame | 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 |
| SpaceInvaders | 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 |

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
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?

2. 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.
soundness: 3
strengths: 1. Scalability: JOWA’s design showcases robust scaling potential, as performance improves with model size without the usual TD-learning instability issues.

2. 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.
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.
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)

Ye, 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.

2. 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.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
questions: How's the model performance when you further scale up the model size, e.g., to 1B?
soundness: 4
strengths: - The proposed JOWA outperforms existing SOTA methods by a large margin. The perform scales up with model sizes

- The joint optimization of world and action models stabilizes large-scale multi-task offline RL training.

- The ablation studies comprehensively study the key design choices of the proposed methods.
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.
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.

- 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?

- 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?

- This work misses the citation for [1,2]

[1] Li et al., Multi-task batch reinforcement learning with metric learning. NeurIPS 2020.

[2] Li et al., Offline reinforcement learning with closed-form policy improvement operators. ICML 2023

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: Please see the weaknesses section.
soundness: 3
strengths: 1. This paper is clearly written and easy to follow.
2. This paper presents sufficient experimental results to demonstrate the validity of its proposed method.
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.
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.
2. Extending the experiments beyond Atari to more complex environments like Kitchen or Meta-World would offer stronger validation of the proposed method's effectiveness.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 4
questions: - Why was the sequence length relative low (8)?
- Could JOWA learn useful policy on new games using non-expert transitions?
- It seems the number for MGDT are very different from the original paper (MGDT achieved 93% IQM HNS on 40M model). Why is this?
- Is the action-part module (CQL) hard to tune in this method?
soundness: 4
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.
- Strong Baseline Performance: It outperforms state-of-the-art methods, showing significant improvements on Atari benchmarks while requiring only 10% of the dataset.
- Adaptability: JOWA demonstrates efficient transfer to new tasks with minimal fine-tuning, highlighting its robustness and generalization capabilities.
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.

JOWA 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.
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.
- 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.
- 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.

---

confidence: 5
contribution: 3
initial_rating: 8
presentation: 3
questions: Q1. Can Theorem 1 be applied for planning of JOWA?
* 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.
* Can you share your thoughts on this? Please correct me in the point if I'm wrong.

Q2. Scaling property in fine-tuning tasks
* Can you provide the fine-tuning results for JOWA-40M, JOWA-70M?
* I think it will be helpful on investigating scaling effects in fine-tuning.

Q3. Extending JOWA to continuous environments
* It would be exciting if JOWA can be also extended to continuous tasks. For example, continous tasks like DMControl are relatively simple but challenging.
* 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.

[1] Mitsuhiko Nakamoto et al., Cal-QL: Calibrated Offline RL Pre-Training for Efficient Online Fine-Tuning, NeurIPS 2023.
soundness: 4
strengths: S1. JOWA significantly outperforms previous multi-task offline RL approaches.
* 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)
* Experimental results of Table 3 shows that JOWA significantly outperforms prior works in fine-tuning (Mean 0.647, IQM 0.647)

S2. JOWA shows favorable scaling property.
* Increase in the number of parameters leads to better performance.
* Compared to other methods, the performance increase is significant.

S3. Ablation studies shows useful insights on JOWA's performance.
* Planning with world model largely contributes the performance.
* Training policy jointly with the world model is important for multi-task agents.

S4. Open-source codes & checkpoints.
* Those resources will be helpful for the community.
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.
weaknesses: W1. Evaluated only in Atari
* For example, continous tasks like DMControl are relatively simple but challenging.
* It will be exciting to see if JOWA can be also applied in continuous tasks.

W2. Discrepancy between Theorem 1 and planning of JOWA
* 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.
* Please correct me in the point if I'm wrong.

Meta-review: This paper introduces a model-based RL agent, JOWA, which trains a single transformer architecture by jointly optimizing the 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 as well, and JOWA can be fine-tuned online and offline as well. 

The reviewers generally liked the paper and the only concern was about training on Atari to plot scaling laws. There's prior work which uses multi-game Atari to evaluate scaling laws, and while I think this is not the end goal for the RL community (and indeed, this paper does only consider 15 Atari games), this work is a good starting point for scaling of model-based RL.

I am also wondering if comparisons to TD-MPC2 could be added to the mix as well (I understand it cannot be applied here directly, but with approximations with action and/or state discretization), since that's a model-based method as well which shows scaling curves, although on different domains. Nonetheless, the paper is of interest and value, and we are accepting this paper.

Decision: Accept (Poster)