GLAM: Global-Local Variation Awareness in Mamba-based World Model

TL;DR

GLAM introduces a dual-module world model that perceives and predicts state variations from global and local perspectives, significantly improving sample efficiency in model-based reinforcement learning, especially on Atari 100k.

Contribution

The paper presents a novel Mamba-based world model with dual reasoning modules, GMamba and LMamba, that better capture state variations to enhance reasoning and learning efficiency.

Findings

Outperforms existing methods on Atari 100k benchmark

Improves sample efficiency in model-based reinforcement learning

Enhances reasoning about environmental variations

Abstract

Mimicking the real interaction trajectory in the inference of the world model has been shown to improve the sample efficiency of model-based reinforcement learning (MBRL) algorithms. Many methods directly use known state sequences for reasoning. However, this approach fails to enhance the quality of reasoning by capturing the subtle variation between states. Much like how humans infer trends in event development from this variation, in this work, we introduce Global-Local variation Awareness Mamba-based world model (GLAM) that improves reasoning quality by perceiving and predicting variation between states. GLAM comprises two Mambabased parallel reasoning modules, GMamba and LMamba, which focus on perceiving variation from global and local perspectives, respectively, during the reasoning process. GMamba focuses on identifying patterns of variation between states in the input sequence and leverages these patterns to enhance the prediction of future state variation. LMamba emphasizes reasoning about unknown information, such as rewards, termination signals, and visual representations, by perceiving variation in adjacent states. By integrating the strengths of the two modules, GLAM accounts for highervalue variation in environmental changes, providing the agent with more efficient imagination-based training. We demonstrate that our method outperforms existing methods in normalized human scores on the Atari 100k benchmark.