[Paper Review] History Compression via Language Models in Reinforcement Learning

Fabian Paischer, Thomas Adler+6 more|arXiv (Cornell University)|May 24, 2022

Reinforcement Learning in Robotics4 citations

TL;DR

This paper proposes HELM (History comprEssion via Language Models), a novel framework that uses a frozen pretrained language Transformer (PLT) as a memory module for history representation in partially observable reinforcement learning. By employing FrozenHopfield—a trainable-free associative memory that maps observations to PLT token embeddings via random projections—it enables sample-efficient policy learning without fine-tuning the PLT, achieving state-of-the-art performance on Minigrid and Procgen environments.

ABSTRACT

In a partially observable Markov decision process (POMDP), an agent typically uses a representation of the past to approximate the underlying MDP. We propose to utilize a frozen Pretrained Language Transformer (PLT) for history representation and compression to improve sample efficiency. To avoid training of the Transformer, we introduce FrozenHopfield, which automatically associates observations with pretrained token embeddings. To form these associations, a modern Hopfield network stores these token embeddings, which are retrieved by queries that are obtained by a random but fixed projection of observations. Our new method, HELM, enables actor-critic network architectures that contain a pretrained language Transformer for history representation as a memory module. Since a representation of the past need not be learned, HELM is much more sample efficient than competitors. On Minigrid and Procgen environments HELM achieves new state-of-the-art results. Our code is available at https://github.com/ml-jku/helm.

Motivation & Objective

• To improve sample efficiency in partially observable reinforcement learning by leveraging pretrained language models for history representation.
• To address the challenge of integrating frozen language models into on-policy RL without fine-tuning or weight updates.
• To design a trainable-free mechanism that maps observations and actions to the input space of a pretrained Transformer.
• To enable actor-critic networks to use compressed, abstract history representations for improved decision-making in POMDPs.
• To achieve state-of-the-art performance on Minigrid and Procgen environments with minimal training overhead.

Proposed method

• Introduces FrozenHopfield, a frozen associative memory that maps observations and actions to pretrained token embeddings using random projections as queries.
• Uses a modern Hopfield network (MHN) to store and retrieve pretrained token embeddings, enabling retrieval without training.
• Employs a fixed projection of observations to generate queries that retrieve relevant token embeddings from the MHN.
• Concatenates the retrieved PLT representation of history with a CNN-embedded current observation to form the input for policy and value networks.
• Uses only shallow multi-layer perceptrons and a CNN as trainable components, leaving the PLT frozen throughout training.
• Applies the framework to actor-critic algorithms like PPO, enabling end-to-end RL training with abstract, compressed history representations.

Experimental results

Research questions

• RQ1 Can a frozen pretrained language model be effectively used for history representation in partially observable RL without fine-tuning?
• RQ2 How can observations from non-language modalities (e.g., pixel observations) be mapped to the input space of a language model in a trainable-free manner?
• RQ3 Does using a frozen PLT for history compression lead to improved sample efficiency compared to standard RNN-based memory mechanisms?
• RQ4 Can a modern Hopfield network serve as a reliable, parameter-free associative memory for retrieving relevant token embeddings from a frozen PLT?
• RQ5 To what extent does HELM outperform existing methods in terms of sample efficiency and final performance on challenging POMDP benchmarks?

Key findings

• HELM achieves state-of-the-art performance on the Minigrid and Procgen environments, outperforming prior methods in sample efficiency and final performance.
• The use of a frozen PLT for history representation leads to significantly improved sample efficiency, as the model does not need to learn the history abstraction from scratch.
• FrozenHopfield enables effective and stable retrieval of relevant token embeddings from the PLT using only random projections and a modern Hopfield network, without any training.
• The retrieval error in the MHN is exponentially small in the separation between patterns, ensuring high-fidelity memory recall.
• Increasing the temperature parameter β in the MHN enhances inter-pattern distance, reducing representation collapse and improving generalization across environments.
• Qualitative analysis shows that similar visual states map to the same or similar tokens in the PLT space, indicating effective abstraction and clustering of semantically similar experiences.