Attention-Based Reward Shaping for Sparse and Delayed Rewards

TL;DR

This paper addresses the challenge of sparsely and delay-rewarded reinforcement learning by introducing Attention-based REward Shaping (ARES), a transformer-based method that derives dense per-step rewards from offline episode data labeled with the final return $G_{ ext{episode}}$. ARES trains offline to predict episodic returns and uses the transformer’s attention to assign credit to individual state-action pairs, producing shaped rewards that can be used with any RL algorithm. The authors demonstrate broad applicability across discrete and continuous environments, with varying data quality (Random vs TrainingExpert) and across baseline algorithms, showing robust improvements in learning under fully delayed rewards and without requiring expert reward models. The results suggest that ARES can significantly accelerate learning in delayed-reward settings and is not limited to goal-based tasks, making it practically impactful for real-world RL where online interaction is costly or infeasible.

Abstract

Sparse and delayed reward functions pose a significant obstacle for real-world Reinforcement Learning (RL) applications. In this work, we propose Attention-based REward Shaping (ARES), a general and robust algorithm which uses a transformer's attention mechanism to generate shaped rewards and create a dense reward function for any environment. ARES requires a set of episodes and their final returns as input. It can be trained entirely offline and is able to generate meaningful shaped rewards even when using small datasets or episodes produced by agents taking random actions. ARES is compatible with any RL algorithm and can handle any level of reward sparsity. In our experiments, we focus on the most challenging case where rewards are fully delayed until the end of each episode. We evaluate ARES across a diverse range of environments, widely used RL algorithms, and baseline methods to assess the effectiveness of the shaped rewards it produces. Our results show that ARES can significantly improve learning in delayed reward settings, enabling RL agents to train in scenarios that would otherwise require impractical amounts of data or even be unlearnable. To our knowledge, ARES is the first approach that works fully offline, remains robust to extreme reward delays and low-quality data, and is not limited to goal-based tasks.

Figures (3)

• Figure 1: ARES Architecture.
• Figure 2: Results on toy environments. Shaded regions are +/- one standard deviation.
• Figure 3: Results on MuJoCo environments. Shaded regions are +/- one standard deviation.