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

TL;DR

GLAM addresses sample efficiency in model-based reinforcement learning by enhancing reasoning about state variation. It introduces two parallel Mamba-based inference modules, LMamba for local variation and GMamba for global variation, to enable richer imagination. The framework integrates both variation perspectives to predict future state distributions and unknown information, with a phased imagined-training schedule to improve data efficiency. On Atari 100k, GLAM achieves state-of-the-art human-normalized scores, especially in games with prominent variation, demonstrating the practical impact of global-local variation awareness for imagination-based training.

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.

Figures (7)

• Figure 1: A single-step demonstration of GLAM inference. GLAM leverages GMamba and LMamba for global and local state variation, respectively. GMamba captures patterns of variation changes within the sequence, while LMamba incorporates variation awareness when inferring unknown information, rather than relying on direct inference alone.
• Figure 2: Overview of the Global-Local variation Awareness Mamba-based world model (GLAM). Based on the known sequence, GLAM infers future information to train the agent in imagination. In each inference, GLAM captures the awareness of global and local variation from long and short sequences and integrates both for prediction.
• Figure 3: Comparison of the training results between GLAM and STORM.
• Figure 4: Ablation study on variable imagine steps (left) and number of layers in the Mamba (right).
• Figure 5: Loss value in the training process of ours w/o G and ours w/o G&L.