2024/02/07 by Michel Ma, Tianwei Ni, Ma, Michel +7 · 3 citations
Chemistry · Engineering · Social Sciences · #Artificial Intelligence (cs.AI) #Chemistry #Economics #Electrical engineering #Engineering #FOS: Computer and information sciences #International Development and Aid #Machine Learning (cs.LG) #Transformer #Voltage
paper · pdf · doi:10.48550/arxiv.2402.05290
published in arXiv (Cornell University) (Cornell University)
openalex publication_date 2024/02/07 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28
A natural approach for reinforcement learning is to predict future rewards by unrolling a neural network world model, and to backpropagate through the resulting computational graph to learn a policy. However, this method often becomes impractical for long horizons since typical world models induce hard-to-optimize loss landscapes. Transformers are known to efficiently propagate gradients over long horizons: could they be the solution to this problem? Surprisingly, we show that commonly-used transformer world models produce circuitous gradient paths, which can be detrimental to long-range policy gradients. To tackle this challenge, we propose a class of world models called Actions World Models (AWMs), designed to provide more direct routes for gradient propagation. We integrate such AWMs into a policy gradient framework that underscores the relationship between network architectures and the policy gradient updates they inherently represent. We demonstrate that AWMs can generate optimization landscapes that are easier to navigate even when compared to those from the simulator itself. This property allows transformer AWMs to produce better policies than competitive baselines in realistic long-horizon tasks.