Improving Token-Based World Models with Parallel Observation Prediction

ICML 24 2024 77.8 method

TLDR

Proposes Parallel Observation Prediction to accelerate token-based world models, achieving 15.4x faster imagination and superhuman Atari performance.

Reasoning

Strengths: addresses a key bottleneck in token-based world models with a novel POP mechanism, demonstrating significant speedup and strong empirical results on Atari 100K. Weaknesses: evaluation is limited to Atari games, and the abstract does not discuss limitations or comparisons to non-token-based approaches.

Read-first score

Read-first score 77.8, weighted from topical fit, citation, graph, method, reproducibility, and recency signals. Original total remains 61.

Topical relevance 42%

87.1

Uses existing LLM keyword relevance scores normalized to 0-100. world model,world simulator,generative world model,interactive world model,video world model,world dynamics prediction,model-based reinforcement learning world model

Reproducibility 25%

Screens links and visible text for paper, code, dataset, artifact, and repository signals. pdf=True; code=True; dataset=False; markers=code,github

Recency 8%

75.1

Uses a gentle age decay so recent papers surface without erasing older foundations. 2024

Methodology quality 25%

Screens visible abstract and analysis fields for experiment, dataset, baseline, metric, and limitation evidence. markers=benchmark,result

Field roles

Reproducibility anchor

Rank sensitivity

Stability: volatile; rank range: 50.

Keyword Scores

world model

model-based reinforcement learning world model

generative world model

world dynamics prediction

world simulator

interactive world model

video world model

Deep Analysis

Innovations

Parallel Observation Prediction (POP) mechanism
Augmenting Retentive Network (RetNet) with a novel forward mode tailored to reinforcement learning
REM (Retentive Environment Model) agent

Methodology

Token-based world models (TBWMs) represent agent experience as language-like sequences of discrete tokens, with each observation as a sub-sequence. The proposed POP mechanism replaces sequential token-by-token generation with parallel observation prediction using a RetNet enhanced with a forward mode designed for reinforcement learning. The REM agent is evaluated on the Atari 100K benchmark.

Key Results

REM achieves a 15.4x speedup in imagination compared to prior TBWMs, attains superhuman performance on 12 out of 26 Atari 100K games, and completes training in under 12 hours.

Resources

Paper Code