Learning Long-Term Reward Redistribution via Randomized Return Decomposition

1 INTRODUCTION . Scaling reinforcement learning ( RL ) algorithms to practical applications has become the focus of numerous recent studies , including resource management ( Mao et al. , 2016 ) , industrial control ( Hein et al. , 2017 ) , drug discovery ( Popova et al. , 2018 ) , and recommendation systems ( Chen et al. , 2018 ) . One of the challenges in these real-world problems is the sparse and delayed environmental rewards . For example , in the molecular structure design problem , the target molecule property can only be evaluated after completing the whole sequence of modification operations ( Zhou et al. , 2019b ) . The sparsity of environmental feedback would complicate the attribution of rewards on agent actions and therefore can hinder the efficiency of learning ( Rahmandad et al. , 2009 ) . In practice , it is a common choice to formulate the RL objective with a meticulously designed reward function instead of the sparse environmental rewards . The design of such a reward function is crucial to the performance of the learned polices . Most standard RL algorithms , such as temporal difference learning and policy gradient methods , prefer dense reward functions that can provide instant feedback for every step of environment transitions . Designing such dense reward functions is not a simple problem even with domain knowledge and human supervision . It has been widely observed in prior works that handcrafted heuristic reward functions may lead to unexpected and undesired behaviors ( Randløv & Alstrøm , 1998 ; Bottou et al. , 2013 ; Andrychowicz et al. , 2017 ) . The agent may find a shortcut solution that formally optimizes the given objective but deviates from the desired polices ( Dewey , 2014 ; Amodei et al. , 2016 ) . The reward designer can hardly anticipate all potential side effects of the designed reward function , which highlights the difficulty of reward engineering . To avoid the unintended behaviors induced by misspecified reward engineering , a common paradigm is considering the reward design as an online problem within the trial-and-error loop of reinforcement learning ( Sorg et al. , 2010 ) . This algorithmic framework contains two components , namely reward modeling and policy optimization . The agent first learns a proxy reward function from the experience data and then optimizes its policy based on the learned per-step rewards . By iterating this procedure and interacting with the environment , the agent is able to continuously refine its reward model so that the learned proxy reward function can better approximate the actual objective given by the environmental feedback . More specifically , this paradigm aims to reshape the sparse and delayed environmental rewards to a dense Markovian reward function while trying to avoid misspecifying the goal of given tasks . In this paper , we propose a novel reward redistribution algorithm based on a classical mechanism called return decomposition ( Arjona-Medina et al. , 2019 ) . Our method is built upon the leastsquares-based return decomposition ( Efroni et al. , 2021 ) whose basic idea is training a regression model that decomposes the trajectory return to the summation of per-step proxy rewards . This paradigm is a promising approach to redistributing sparse environmental feedback . Our proposed algorithm , randomized return decomposition ( RRD ) , establish a surrogate optimization of return decomposition to improve the scalability in long-horizon tasks . In this surrogate problem , the reward model is trained to predict the episodic return from a random subsequence of the agent trajectory , i.e. , we conduct a structural constraint that the learned proxy rewards can approximately reconstruct environmental trajectory return from a small subset of state-action pairs . This design enables us to conduct return decomposition effectively by mini-batch training . Our analysis shows that our surrogate loss function is an upper bound of the original loss of deterministic return decomposition , which gives a theoretical interpretation of this randomized implementation . We also present how the surrogate gap can be controlled and draw connections to another method called uniform reward redistribution . In experiments , we demonstrate substantial improvement of our proposed approach over baseline algorithms on a suite of MuJoCo benchmark tasks with episodic rewards . 2 BACKGROUND . 2.1 EPISODIC REINFORCEMENT LEARNING WITH TRAJECTORY FEEDBACK . In standard reinforcement learning settings , the environment model is usually formulated by a Markov decision process ( MDP ; Bellman , 1957 ) , defined as a tuple M = 〈S , A , P , R , µ〉 , where S and A denote the spaces of environment states and agent actions . P ( s′|s , a ) and R ( s , a ) denote the unknown environment transition and reward functions . µ denotes the initial state distribution . The goal of reinforcement learning is to find a policy π : S → A maximizing cumulative rewards . More specifically , a common objective is maximizing infinite-horizon discounted rewards based on a pre-defined discount factor γ as follows : ( standard objective ) J ( π ) = E [ ∞∑ t=0 γtR ( st , π ( st ) ) ∣∣∣∣∣ s0 ∼ µ , st+1 ∼ P ( · | st , π ( st ) ) ] . ( 1 ) In this paper , we consider the episodic reinforcement learning setting with trajectory feedback , in which the agent can only obtain one reward feedback at the end of each trajectory . Let τ denote an agent trajectory that contains all experienced states and behaved actions within an episode . We assume all trajectories terminate in finite steps . The episodic reward function Rep ( τ ) is defined on the trajectory space , which represents the overall performance of trajectory τ . The goal of episodic reinforcement learning is to maximize the expected trajectory return : ( episodic objective ) Jep ( π ) = E [ Rep ( τ ) ∣∣∣∣ s0 ∼ µ , at = π ( st ) , τ = 〈s0 , a0 , s1 , · · · , sT 〉 ] . ( 2 ) In general , the episodic-reward setting is a particular form of partially observable Markov decision processes ( POMDPs ) where the reward function is non-Markovian . The worst case may require the agent to enumerate the entire exponential-size trajectory space for recovering the episodic reward function . In practical problems , the episodic environmental feedback usually has structured representations . A common structural assumption is the existence of an underlying Markovian reward function R̂ ( s , a ) that approximates the episodic reward Rep ( τ ) by a sum-form decomposition , ( sum-decomposable episodic reward ) Rep ( τ ) ≈ R̂ep ( τ ) = T−1∑ t=0 R̂ ( st , at ) . ( 3 ) This structure is commonly considered by both theoretical ( Efroni et al. , 2021 ) and empirical studies ( Liu et al. , 2019 ; Raposo et al. , 2021 ) on long-horizon episodic rewards . It models the situations where the agent objective is measured by some metric with additivity properties , e.g. , the distance of robot running , the time cost of navigation , or the number of products produced in a time interval . 2.2 REWARD REDISTRIBUTION . The goal of reward redistribution is constructing a proxy reward function R̂ ( st , at ) that transforms the episodic-reward problem stated in Eq . ( 2 ) to a standard dense-reward setting . By replacing environmental rewards with such a Markovian proxy reward function R̂ ( st , at ) , the agent can be trained to optimize the discounted objective in Eq . ( 1 ) using any standard RL algorithms . Formally , the proxy rewards R̂ ( st , at ) form a sum-decomposable reward function R̂ep ( τ ) = ∑T−1 t=0 R̂ ( st , at ) that is expected to have high correlation to the environmental reward Rep ( τ ) . Here , we introduce two branches of existing reward redistribution methods , return decomposition and uniform reward redistribution , which are the most related to our proposed approach . We defer the discussions of other related work to section 5 . Return Decomposition . The idea of return decomposition is training a reward model that predicts the trajectory return with a given state-action sequence ( Arjona-Medina et al. , 2019 ) . In this paper , without further specification , we focus on the least-squares-based implementation of return decomposition ( Efroni et al. , 2021 ) . The reward redistribution is given by the learned reward model , i.e. , decomposing the environmental episodic reward Rep ( τ ) to a Markovian proxy reward function R̂ ( s , a ) . In practice , the reward modeling is formulated by optimizing the following loss function : LRD ( θ ) = E τ∼D [ ( Rep ( τ ) − T−1∑ t=0 R̂θ ( st , at ) ) 2 ] , ( 4 ) where R̂θ denotes the parameterized proxy reward function , θ denotes the parameters of the learned reward model , and D denotes the experience dataset collected by the agent . Assuming the sumdecomposable structure stated in Eq . ( 3 ) , R̂θ ( s , a ) is expected to asymptotically concentrate near the ground-truth underlying rewards R̂ ( s , a ) when Eq . ( 4 ) is properly optimized ( Efroni et al. , 2021 ) . One limitation of the least-squares-based return decomposition method specified by Eq . ( 4 ) is its scalability in terms of the computation costs . Note that the trajectory-wise episodic reward is the only environmental supervision for reward modeling . Computing the loss function LRD ( θ ) with a single episodic reward label requires to enumerate all state-action pairs along the whole trajectory . This computation procedure can be expensive in numerous situations , e.g. , when the task horizon T is quite long , or the state space S is high-dimensional . To address this practical barrier , recent works focus on designing reward redistribution mechanisms that can be easily integrated in complex tasks . We will discuss the implementation subtlety of existing methods in section 4 . Uniform Reward Redistribution . To pursue a simple but effective reward redistribution mechanism , IRCR ( Gangwani et al. , 2020 ) considers uniform reward redistribution which assumes all state-action pairs equally contribute to the return value . It is designed to redistribute rewards in the absence of any prior structure or information . More specifically , the proxy reward R̂IRCR ( s , a ) is computed by averaging episodic return values over all experienced trajectories containing ( s , a ) , R̂IRCR ( s , a ) = E τ∼D [ Rep ( τ ) | ( s , a ) ∈ τ ] . ( 5 ) In this paper , we will introduce a novel reward redistribution mechanism that bridges between return decomposition and uniform reward redistribution . 3 REWARD REDISTRIBUTION VIA RANDOMIZED RETURN DECOMPOSITION . In this section , we introduce our approach , randomized return decomposition ( RRD ) , which sets up a surrogate optimization problem of the least-squares-based return decomposition . The proposed surrogate objective allows us to conduct return decomposition on short subsequences of agent trajectories , which is scalable in long-horizon tasks . We provide analyses to characterize the algorithmic property of our surrogate objective function and discuss connections to existing methods . 3.1 RANDOMIZED RETURN DECOMPOSITION WITH MONTE-CARLO RETURN ESTIMATION . One practical barrier to apply least-squares-based return decomposition methods in long-horizon tasks is the computation costs of the regression loss in Eq . ( 4 ) , i.e. , it requires to enumerate all state- action pairs within the agent trajectory . To resolve this issue , we consider a randomized method that uses a Monte-Carlo estimator to compute the predicted episodic return R̂ep , θ ( τ ) as follows : R̂ep , θ ( τ ) = T−1∑ t=0 R̂θ ( st , at ) ︸ ︷︷ ︸ Deterministic Computation = E I∼ρT ( · ) [ T |I| ∑ t∈I R̂θ ( st , at ) ] ≈ T |I| ∑ t∈I R̂θ ( st , at ) ︸ ︷︷ ︸ Monte-Carlo Estimation , ( 6 ) where I denotes a subset of indices . ρT ( · ) denotes an unbiased sampling distribution where each index t has the same probability to be included in I . In this paper , without further specification , ρT ( · ) is constructed by uniformly sampling K distinct indices . ρT ( · ) = Uniform ( { I ⊆ ZT : |I| = K } ) , ( 7 ) whereK is a hyper-parameter . In this sampling distribution , each timestep t has the same probability to be covered by the sampled subsequence I ∼ ρT ( · ) so that it gives an unbiased Monte-Carlo estimation of the episodic summation R̂ep , θ ( τ ) . Randomized Return Decomposition . Based on the idea of using Monte-Carlo estimation shown in Eq . ( 6 ) , we introduce our approach , randomized return decomposition ( RRD ) , to improve the scalability of least-squares-based reward redistribution methods . The objective function of our approach is formulated by the randomized return decomposition loss LRand-RD ( θ ) stated in Eq . ( 8 ) , in which the parameterized proxy reward function R̂θ is trained to predict the episodic return Rep ( τ ) given a random subsequence of the agent trajectory . In other words , we integrate the Monte-Carlo estimator ( see Eq . ( 6 ) ) into the return decomposition loss to obtain the following surrogate loss function : LRand-RD ( θ ) = E τ∼D [ E I∼ρT ( · ) [ ( Rep ( τ ) − T |I| ∑ t∈I R̂θ ( st , at ) ) 2 ] ] . ( 8 ) In practice , the loss function LRand-RD ( θ ) can be estimated by sampling a mini-batch of trajectory subsequences instead of computing R̂θ ( st , at ) for the whole agent trajectory , and thus the implementation of randomized return decomposition is adaptive and flexible in long-horizon tasks .

Summary: This paper targets how to decompose delayed reward signals obtained at the end of a trajectory to a more immediate form of reward feedback. The authors introduce an algorithm, "Reward Distribution vis Randomized Return Decomposition" that serves as a proxy reward function for episodic reinforcement learning. The authors showcase extensive experiments showcasing improved performance over relevant baselines on common Mujoco benchmarks. The main contribution is to convert long horizon delayed reward problems to more short length sequences that can be trained upon with mini-batch gradient descent.