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: The paper addresses the problem of credit assignment in delayed reward setting. It does this by providing a new mechanism for reward redistribution. The authors claim that the new mechanism makes reward redistribution more scalable. The authors predict the return of a trajectory by using only random sub-sequences in the trajectory. Then the prediction model is used for assigning reward for state-action pairs.