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
Primary area: reinforcement learning
Keywords: reinforcement learning, offline reinforcement learning, world model

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.

Decision: accept (poster)
Reviewer scores: 5, 8, 6, 6, 8

Reviews and discussion:
<thread>
creation date:
1729349959205

thread_type:
official_review

<thread_object>
creator: reviewer_xu9g

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

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.

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.

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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730278910879

thread_type:
official_review

<thread_object>
creator: reviewer_plgr

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

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.

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.

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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730557155615

thread_type:
official_review

<thread_object>
creator: reviewer_pkew

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

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.

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.

questions:
Please see the weaknesses section.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730715365944

thread_type:
official_review

<thread_object>
creator: reviewer_7vxq

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

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.

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

questions:
How's the model performance when you further scale up the model size, e.g., to 1B?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730848532082

thread_type:
official_review

<thread_object>
creator: reviewer_zgrd

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

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.

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.

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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732543860476

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Update manuscript

comment:
We thank reviewers for all the valuable feedback. We address all the reviewers' comments below and have incorporated all feedback in the revised manuscript  using blue font. Specifically, we have updated the following:
1. Add several relevant works in section 2 (Reviewer zGRd, Reviewer 7Vxq).
2. Employ FICC as a model-based baseline for main experiments in section 5 (Reviewer zGRd).
3. Add JOWA's scaling results and conclusions in the fine-tuning experiments, i.e., section 5.5 (Reviewer XU9g).
4. Add a comparison with MCTS in section 5.6 to emphasize the efficiency of our planning algorithm (Reviewer zGRd).
5. Add a description of the scaling inconsistency in the limitations (Reviewer 7Vxq).
6. Add additional experiments in Appendix F, including details of MCTS (Reviewer zGRd), fine-tuning results using non-expert data (Reviewer pLGR), and explanations of the emergent behaviors (Reviewer 7Vxq).

We sincerely aspire that our detailed rebuttal will dispel any uncertainties or misunderstandings which reviewers may have raised regarding our manuscript, thus contributing positively to the final ratings of this work. If any additional experiments are needed to further demonstrate the potential of JOWA, we will do our utmost to supplement the relevant experiments during the valuable discussion period.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733303953829

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of Rebuttal

comment:
We are deeply grateful to all reviewers for their thorough evaluation of our work and their constructive feedback. We particularly appreciate their recognition of our paper's key strengths:

1. Strong Baseline Performance: Reviewers 7Vxq, pLGR, and XU9g
2. Scalability: Reviewers zGRd, 7Vxq, pLGR, and XU9g
3. Adaptability: Reviewer pLGR
4. Comprehensive Ablation Studies: Reviewers zGRd, 7Vxq, and pKew
5. Open-Source codes: Reviewer XU9g

Through productive discussions during the rebuttal period, we have made substantial improvements to our paper to address the reviewers' concerns:

1. Added additional relevant works (Reviewers zGRd, 7Vxq)
2. Added a model-based baseline for main experiments (Reviewer zGRd)
3. Added JOWA's scaling results and conclusions in fine-tuning experiments (Reviewer XU9g)
4. Added a comparison between our planning algorithm and MCTS (Reviewer zGRd)
5. Added a discussion of scaling inconsistency in the limitations section (Reviewer 7Vxq)
6. Added fine-tuning results using non-expert data (Reviewer pLGR)
7. Added the explanation of emergent behaviors (Reviewer 7Vxq)

After rebuttal, the main remaining concern from reviewers is that Atari might not be an ideal benchmark for evaluating RL's scaling law (Reviewer zGRd). We will carefully consider the reviewers' suggestions and design experiments based on their comments in our subsequent series of work, such as expanding to continuous action spaces.

We extend our sincere gratitude to the reviewers for their valuable insights, and to the Area Chair, Senior Area Chair, and Program Committee for their dedication in organizing this review process!
</content>
</thread_object>

</thread>


Metareview:
metareview:
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.