[Paper Review] Beyond Exponentially Discounted Sum: Automatic Learning of Return Function

Yufei Wang, Qiwei Ye+1 more|arXiv (Cornell University)|May 28, 2019

Reinforcement Learning in Robotics36 references4 citations

TL;DR

This paper proposes a meta-learning approach to automatically discover optimal return functions in reinforcement learning, moving beyond the standard exponentially discounted sum. By modeling returns as learnable linear combinations of trajectory rewards using a transformer-based coefficient predictor, the method accelerates learning in sparse-reward environments like mazes and Atari games, with empirical results showing significant performance gains in delayed-reward settings.

ABSTRACT

In reinforcement learning, Return, which is the weighted accumulated future rewards, and Value, which is the expected return, serve as the objective that guides the learning of the policy. In classic RL, return is defined as the exponentially discounted sum of future rewards. One key insight is that there could be many feasible ways to define the form of the return function (and thus the value), from which the same optimal policy can be derived, yet these different forms might render dramatically different speeds of learning this policy. In this paper, we research how to modify the form of the return function to enhance the learning towards the optimal policy. We propose to use a general mathematical form for return function, and employ meta-learning to learn the optimal return function in an end-to-end manner. We test our methods on a specially designed maze environment and several Atari games, and our experimental results clearly indicate the advantages of automatically learning optimal return functions in reinforcement learning.

Motivation & Objective

• To address the limitation of fixed return functions in reinforcement learning, especially the slow learning caused by exponential discounting in sparse-reward scenarios.
• To investigate whether learning a more effective return function form—beyond exponential discounting—can accelerate policy learning while preserving the same optimal policy.
• To develop an end-to-end, differentiable method that learns the structure of the return function using meta-gradient optimization.
• To demonstrate that trajectory-aware return functions, which redistribute credit more effectively, lead to faster convergence in challenging RL environments.

Proposed method

• Proposes a general return function form: $ G^{g}(s_t, a_t) = \sum_{t=0}^{T} \beta_t r_t $, where coefficients $ \beta_t $ are dynamically predicted from the full trajectory.
• Uses a transformer-based network to compute the coefficients $ \beta_t $, enabling relational reasoning over states and rewards in a trajectory.
• Employs meta-learning (specifically, meta-gradient descent) to optimize the return function parameters so that the downstream policy learns faster.
• Integrates the learned return function into an actor-critic framework (A2C), with the value function trained to predict the meta-learned return.
• Treats the return function as a differentiable module, allowing backpropagation through the return computation to update the coefficient network.
• Searches over hyperparameters (learning rate and trajectory length) to find optimal configurations for each environment.

Experimental results

Research questions

• RQ1 Can a return function that redistributes credit more effectively than exponential discounting lead to faster policy learning in reinforcement learning?
• RQ2 Is it possible to automatically learn the optimal form of the return function in an end-to-end manner using meta-learning?
• RQ3 How does a trajectory-aware, non-exponential return function affect learning speed in environments with sparse or delayed rewards?
• RQ4 Does the proposed method generalize across environments with varying reward sparsity, such as mazes and Atari games?

Key findings

• In the specially designed maze environment, the proposed RGM method achieved significantly faster learning than both vanilla A2C and standard exponential discounting, particularly in the early training phase.
• The learned return function in the maze environment concentrated nearly all weight on the final delayed reward, enabling near-lossless backpropagation of the learning signal to earlier states.
• On Atari games with delayed rewards (e.g., Frostbite, Bowling), A2C + RGM showed substantial performance improvements over vanilla A2C and LIRPG, demonstrating faster convergence.
• For games with dense immediate rewards (e.g., Space Invaders, Breakout), the method maintained competitive performance, indicating robustness across reward sparsity levels.
• The method outperformed baseline A2C in 7 out of 10 tested Atari games, with the largest gains observed in environments where the original return function struggled due to delayed credit assignment.
• Visualizations confirmed that the RGM learned to assign high coefficients to the correct path's final reward, while the exponential baseline decayed this signal exponentially over time.