CMU-CS-26-114
Computer Science Department
School of Computer Science, Carnegie Mellon University



CMU-CS-26-114

Attention Over The Past for Data Efficiency in RL

Allen Zheng

M.S. Thesis

May 2026

CMU-CS-26-114.pdf


Keywords: Reinforcement Learning, Policy Gradient Methods, Advantage Estimation, Sample Efficiency

We introduce a framework that incorporates attention over past states into reinforcement learning (RL) in two complementary ways. First, we use an attention mechanism over a trajectory buffer of previously visited states to construct a history-aware critic, replacing the standard neural network critic with an estimate computed as an attention-weighted average over stored values. Second, we replace GAE, which is computed through a single trajectory, with an attention-weighted advantage that, in addition to stepping forward through time, also steps according to similar states. States more similar to the current one contribute more to the return estimate, providing a smooth, similarity-weighted alternative to the sequential rollout. Together, these two mechanisms reduce variance by pooling signals across similar states rather than relying on a single trajectory.

38 pages

Thesis Committee:
Geoffrey J. Gordon (Chair)
Jeff Schneider

Jignesh Patel, Interim Head, Computer Science Department
Martial Hebert, Dean, School of Computer Science


Return to: SCS Technical Report Collection
School of Computer Science

This page maintained by reports@cs.cmu.edu