Reinforcement Learning from Bagged Reward
Authors: Yuting Tang
Xin-Qiang Cai
Yao-Xiang Ding
Qiyu Wu
Guoqing Liu
Masashi Sugiyama
Venue: Accepted by TMLR ()

Abstract: In Reinforcement Learning (RL), it is commonly assumed that an immediate reward signal is generated for each action taken by the agent, helping the agent maximize cumulative rewards to obtain the optimal policy. However, in many real-world scenarios, designing immediate reward signals is difficult; instead, agents receive a single reward that is contingent upon a partial sequence or a complete trajectory. In this work, we define this challenging problem as RL from Bagged Reward (RLBR), where sequences of data are treated as bags with non-Markovian bagged rewards, leading to the formulation of Bagged Reward Markov Decision Processes (BRMDPs). Theoretically, we demonstrate that RLBR can be addressed by solving a standard MDP with properly redistributed bagged rewards allocated to each instance within a bag. Empirically, we find that reward redistribution becomes more challenging as the bag length increases, due to reduced informational granularity. Existing reward redistribution methods are insufficient to address these challenges. Therefore, we propose a novel reward redistribution method equipped with a bidirectional attention mechanism, enabling the accurate interpretation of contextual nuances and temporal dependencies within each bag. We experimentally demonstrate that the proposed method consistently outperforms existing approaches.

Published in Transactions on Machine Learning Research (04/2025)
Reinforcement Learning from Bagged Reward
Yuting Tang'
lunggmms.k..lllkgg..ac.
The University of Tokyo 8 RIKEN Center for AIP
Tokyo, Jopan
Xin-Qiang Cai*
RIKEN Center for AIP
Tokyo, Japan
Yao-Xiang Ding
dingya. gm @gmail com
State Key Lab for CAD & CG, Zhejiang University Hangzhou, China
Qiyu Wu
giguu@@oogss..u-ttkkgo..coo
The University of Tokyo
Tokyo, Japan
Guoqing Liu
guoginglia&mmccroooft.com
Microsoft Research A14Science
Beijöng, China
Masashi Sugiyama
sugi@k..utokyyo.a..
RIKEN Center for ATP & The University of Tokyo Tokyo, Japan
Reviewed on OpenReview: https: // openreview. net/ forum? id= bXUspBbZDA
Abstract
In Reinforcement Learning (RL), it is commonly assumed that an immediate reward signal is generated for each action taken by the agent, helping the agent maximize cumulative rewards to obtain the optimal policy. However, in many real-world scenarios, designing immediate reward signals difficult; instead, agents receive a single reward that contingent upon partial sequence a complete trajectory. In this work, define this challenging problem as RL from Bagged Reward (RLBR), where sequences of data are treated bags with non- Markovian bagged rewards, leading to the formulation of Bagged Reward Markov Decision Processes (BRMDPs). Theoretically, we demonstrate that RLBR can be addressed by solv- ing standard MDP with properly redistributed bagged rewards allocated to each instance within a bag. Empirically, we find that reward redistribution becomes more challenging as the bag length increases, due to reduced informational gramularity. Existing reward redistri- bution methods are insufficient. to address these challenges. Therefore, propose novel reward redistribution method equipped with bidirectional attention mechanism, enabling the accurate interpretation of contextual nuances and temporal dependencies within each bag. We experimentally demonstrate that the proposed method consistently outperforms existing approaches. Codes are available in
com/Tang-YYt.nn//RBB
"Equal contribution.

---

Published in Transactions on Machine Learning Research (04/2025)
1 Introduction
Reinforcement Learning (RL) has achieved remarkable success in various domains, including autonomous driving (Kiran ct al., 2021), healtheare (Yu et al. 2021] complex game playing (Silver et al. 2016 Wurman et al. 2022), and financial trading (Yang et al., 2020 One common and essential assumption for most
RL algorithms is the availability of immediate reward feedback at each time step of the decision-makiin making process. However, designing such immediate reward feedback is quite difficult and may mislead the policy learning in many real-world applications (Kwon et al.) 2023: Lee ct al. 2023). Recognizing this gap, numerous studies (Watkins 1989 Ke et al. 2018: Arjona-Medina et al. 2019 Gangwani et al. 2020: Ren et, al.) 2021;
Zhang et al.) 2023 have explored the concept of RL with Trajectory Feedback (RLTF), primarily focusing
trajectory feedback where rewards are allocated at the end of sequence.
However, real-world applications often feature complex non-immediate reward structures that traditional RLTF cannot fully capture. For instance, in autonomous driving (see Fi rewards are typically linked to completing specific tasks by sequences of actions rather than individual actions or just the final objective (Early et al., 2022; Gaon & Brafman 2020). Similarly, in mobile health interventions (Gao et al., 2024). : the reward is defined by sustained daily commitment to being active throughout the intervention period, with multiple activity suggestions each day influencing this outcome. In such scenarios, providing rewards for every action is impractical, while focusing only on the final goal overlooks important aspects of the decision- making process. Previous studies have explored learning desirable policies using immediate rewards (ewardsSSutton
& Barto 1998 Vinyals et al.) 2019) or trajectory-baste feodback (Arjona-Medina ct al. 2019
Gangwani
et al.
2020 Ren et al. 2021] | are not ideal methods for such scenarios.
To address these challenges, we introduce RL from Bogged Rewards (RLBR), a framework that better aligns with real-world scenarios by considering the cumulative effect of scries of instances. In RLBR, sequences of instances are defined as bags, each associated with a bugged reward. This framework encompasses both the traditional RL setting, where each bag contains only a single instance, and the trajectory feedback setting, where a bag spans the entire trajectory, as special cases. Furthermore, RLBR offers the potential to reduce the labeling workload by lessening the frequency of reward annotations. However, this benefit is balanced by increased learning complexity due to the reduced gramularity of information.
In the RLBR framework, our focus is on leveraging bagged reward information to discern the significance of each instance within bag and to understand the relationships among different bags. The challenge lies in accurately interpreting the contextual nuances within individual bags, the instances within bag would be time-dependent on cach other and their contributions to the bagged reward vary. Given the importance of contexts in RLBR, we turn to the bidirectional allention mechanism (Seo el al. 2016 Vaswani et al. 2017 Devlin et al. 2019). renowned for its effectiveness in contextual understanding, especially for time-dependent reward model, leveraging the bidiroctional attention
data. Specifically, we propose a Transformer-basee mechanism to adeptly interpret contexts within a bag and allocate rewards each instance accurately. This model can be utilized to enhance general RL algorithms, such as Soft Actor-Critic (SAC) (Haarnoja et al., 2018), for environments with bagged rewards.
In summary, our contributions are threefold:
1. We introduce the RLBR framework as a general problem sctting and formalize it using an extension of traditional Markov Decision Processes (MDPs), which we define as Bagged Reward MDPs (BRMDPs). This flexible reward structure bridges instance-level and trajectory-level rewards by supporting variable- length reward bags, making it suitable for real-world scenarios (Section 3)-
2. We establish a theoretical connection between MDPs after reward redistribution and the original BR- MDPs, providing a solid foundation for policy learning under bagged rewards. Building on this, we proposc Transformer-basee reward model with a bidirectional attention mechanism that effectively cap- tures temporal and contextual dependencies, combined with an alternating optimization algorithm to jointly enhance reward redistribution and policy learning (Section 4
3. Extensive experiments show that our method not. only outperforms existing approaches but also maintains robust performance as bag length increases, highlighting its ability to mimic ground truth reward distribu- tions while cxhibiting contextual understanding and adaptability to environmental dynamics (Soction 5
2

---

Published in Transactions on Machine Learning Research (04/2025)
Score 70
Score
Score 100
Check
20
Point
P
Check Point
Start +
Check
Point
L Reward Bag
End
Reward Bag
Reward Bag
Figure 1: An illustration of the reward bag structure in an autonomous driving scenario. The driving trajectory is divided into multiple segments, or bags, each representing sequence of actions that is evaluated and assigned cumulative reward by an external evaluator.
2 Related Works
With the growing attention on RLTF, where trajectory feedback is often considered the sum of rewards over the entire trajectory, new methods specifically designed for RLTF have emerged. Return Decomposition for Delayed Rewards (RUDDER) (Arjona-Medina et al. 2019) used a return-equivalent formulation for precise credit assignment, with advancements incorporating expert demonstrations and language models (Liu et al. Widrich al., 2021 Patil al., 2022). Both Iterative Relative Credit Refinement (IRCR) (Gangwani 2020g and Randomized Return Decomposition (RRD) (Ren et al. 2021) assumed that each state- action pair contributes equally to the trajectory feodback. Specifically, IRCR, presented uniform reward redistribution model, while RRD proposed novel upper bound for return-equivalent assumptions, inte- grating return decomposition with uniform reward redistribution. Additionally, Han et al. (2022) proposed modifying RI. algorithms to incorporate sequence-level information, enabling agents to learn from broader structures and long-term outcomes. On the other hand, some studies on sparse rewards, such as reward shaping (Ng et al.] 1999 Tambwekar al. 2019 Hu et al. 2020) and intrinsic reward methods (Sorg et al., 2010) Pathak et al., 20177 Zheng et al., 2018 2021), could be directly applied to RLTF. However, these approaches primarily focused addressing the exploration-cxploitatii trade-off in sparse reward scttings, which made them less ideal for solving RLTF problems. A detailed discussion on the sparse reward settings and methods is provided in Appendix
2019
et
Despite recent advances, methods that assume equal contribution from cach state-action pair to trajectory feedback overlook the varying importance of individual instances. On the other hand, methods that consider sequence information struggled to capture effective information from long trajectories, resulting in suboptimal training outcomes, as evidenced by prior studies (Ren al., 2021 Zhang al. 2023). Building on previous methods, adopt reward redistribution learning strategy to enhance policy learning in the context of reward bags. The proposed method aims to extract information from bag-level rewards, taking into account the contribution of cach instance, ultimately learning a reliable and well-performing policy. Experimental results (see Sontion 5) demonstrate that the proposed method can accurately capture information from long sequences, and the rich information in the bags further enhances its effectiveness.
3
Reinforcement Learning from Bagged Reward
In this section, we first provide preliminaries on RL with immediate rewards and trajectory leedback, and then formulate the RLBR problem with an extension of the traditional MDPs, the BRMDPs.
3.1 Preliminaries
In traditional RL problems, which are modeled using MDPs, each state-action pair is promptly associated with reward (Sutton & Barto 1998). This paradigm is encapsulated in tuple M (S, P, r.), with
3

---

Published in Transactions on Machine Learning Research (04/2025)
8 9 and A M as finite sets of states and actions, P p(s'|s, a) state transition probability function to state s' when action is taken at state 8, r(s, a) as the immediate reward function when action is taken at state 8, and μ p(s) as the initial state distribution. The primary objective in this framework is to discover a policy π = p(als), which is a conditional probability of taking action a given state s, that maximizes the cunulative sum of rewards over a horizon length T:
T-1 Σ
J(x) = ET P.
r(tt, ae)
(1)
Given that associating each state-action pair with a reward is challenging in real-world scenarios and in- volves significant labeling costs, RLTF, also known as episodic or delayed rewards in some works (Watkins) Zhang et al. 2023), has become increasingly prominent in many ap- plications (Agarwal et al. 2021 Chen et al. 2024). Distinct. from traditional RL, RLTF offers only one feedback after complete trajectory |Arjona-Meding et al. 2019 Zhang et al. 2023 A trajectory {(50, un), (51, a.)...., (0r)...,,,ss1..ppp..) (rT=)} consists of T state-action pairs, with a cutnulative reward (ama(T) that is the sum of latent immodiate rewards Σt= r(st, at), observable only at the end. Denoting by T(x) the distribution of trajectories induced by π. P. p4, the learning objective in RLTF is to maximize the expected rajeectory-bssed reward:
1989. Arjona-Medina et al. 2019:
Jiray(T) (T()) [T)r
(2)
3.2 Problem Formulation
We formulate the RLBR problem as an extension of the traditional MDP, referred to as the BRMDP, which features reward granularity that lies between the two aforementioned settings. We begin by defining the con- cept of bags, which are sub-picces of complete trajectorics. A trajectory divided into several neighboring bags, and bag of size nu, which starts from time i, is defined Br,,; = {(s;;,a),,.. ai), ...,. (si4n, 1114,,,11 1, ai+n; 1)), S i n, - 1 < T - 1. Afterward, define the BRMDP to navigate the complexities of the aggregated non-Markovian reward signals:
Definition 1 (BRMDP). A BRMDP is defined by the tuple (8, P, R, where
S and A are finite sets of states and actions.
P is the state transition probubility function.
R denotes the bagged reward function, which defines the reward over bag: R R(B....)
represents the initial state distribution.
In the RLBR framework, a bag Bin metaphorically a unified reward unit for a contiguous sequence of state-action pairs. Importantly, we do not impose any specific structural assumption on the bagged reward itself. It can be formed through various aggregation mechanisms beyond simple summation of instance-level reward, depending
the task and environment. A trajectory is composite of a set of bags, denoted B,, which ensures that each trajectory includes at least one reward bag. We further assume bag partition function defined by the environment: G : T + B,, which is task-dependent function for generating bags given an input trajectory. Accordingly, the objective is to learn a policy æ that maximizes the expected cumulative leg-level reward:
JB(T) = BT(r)
Σ R(B) g BEB.
Notably, the RLBR framework subsumes several existing paradigms as special cases. When each bag consists 1), a single instance (n, for 1, 1), RLBR simplifics to traditional RL. Converscly, when a single bag covers the entire trajectory (no T), RLBR reduces to RLTF. This adaptability highlights the flexibility of the RLBR framework in accommodating various reward structure scenarios.
4

---

Published in Transactions on Machine Learning Research (04/2025)
4 Reward Redistribution for RLBR
This section delves into the proposed reward redistribution method for RLBR. We aim for the reward redis- tribution to meet the following criteria: first, the policies learned from the MDP after reward redistribution should be consistent with those learned from the BRMDP; second, the redistribution process should reflect each instance's 8 contribution based on the sequence information. The subsequent subsections are structured to first theoretically prove the rationality of reward redistribution, followed by proposing the use of bidi- rectional attention mechanism for this redistribution. Finally, we outline comprehensive algorithm for developing efficient policies within the BRMDP framework.
4.1 Equivalence of Optimal Policies in the BRMDP and the MDP after Reward Redistribution
In the BRMDP, agents do not have direct access to the immediate rewards associated with each in- stance. To address this, we adopt common approach used in trajectory feedback scenarios, modeling the cumulative sum form as the reward redistribution strategy to convert bagged rewards into instance- level rewards (Ren et al. 2021 Arjona-Medina et al. 2019 Zhang et al. 2023 Specifically, we define R(B(...)= = - f(s,, a,), where f(s,, az) represents the redistributed reward. However, a critical ques- tion arises: How can we ensure that the policy learned with redistributed rewards performs well in the original BRMDP? To address this, we provide the following thcorem:
Theorem 1. Consider a BRMDP where only the sequence-level bagged reward available. Suppose the bagged reward is redistributed across instances while preserving its total sum within each bag. Under this assumplion, the set of optimal policies II learned in the resulting MDP is identical to the set of oplimal policies MB in the BRMDP; meaning that II ll.. [B.
rewards from all the bags along the trajectory: Proof. Over complete trajectory т, the cumulative reward in BRMDP can be expressed as the sum of the
it Σ
T-1
Σ R(B) ((B
Σ
= Σ f(st,
BEB,
VEI.
İEI.
1=1
where I, denotes the set of initial timestep indices for B B,..
The policy optimization objective in BRMDP, JB(x), aims to maximize the expected sum of bagged rewards along the trajectory from t = O, which is equivalent to maximizing the cumulative reward in redistributed reward MDP, (TT::
Σ RBD) BEB.
Ex P. Σ Fln ((x)
JB(T) ET(m)
Given that the expected cumulative rewards for any policy in the original BRMDP and redistributed reward MDP frameworks are equivalent, under the condition of infinite exploration or exhaustive sampling within the state-action space, the sets of optimal policies for each framework also coincide, implying that IT IB.
Under this theorem, directly using bagged rewards or applying an average reward redistribution strategy can, in theory, lead to an optimal policy. However, directly using bagged rewards introduces delayed feedback, which slows down policy convergence, demonstrated in Arjona-Medina et al. (2019). Meanwhile, average reward redistribution assumes equal contributions from all instances within bag, potentially diluting critical decision points and consequently affecting policy learning (Gangwani et al. 2020 Ren et al. (021). Therefore, an effective reward redistribution method should capture the actual contributions of individual stale-action pairs within each bag. By assigning rewards in a way that reflects their true impact on bagged rewards, we can enhance credit assignment accuracy and improve the quality of policy optimization. By this thcorem, we ensure that optimizing a policy on an MDP after rewards redistribution is equivalent to optimizing on the original BRMDP. Below, we introduce our specific reward redistribution approach.
5

---

Published in Transactions on Machine Learning Research (04/2025)
Reward Prediction Puni 5 6, ((0;Bix)= (2i0) 7 R R(B(ai))
Bidirectional Attention Layer
Nost State Prediction (5,, Sc:5 L, (0)0 S141) = 18844 Seas
91=1
Causal Transformer
Sequence of Input Data
The illustration of the Reward Bag Transformer (RBT) architecture. The Causal Transformer is used for reward representation by processing sequences of input data consisting of state-action pairs. The bidirectional attention layer is used for reward redistribution utilizing the outputs of the Causal Transformer to predict instance-level rewards. The next state prediction helps the model understand the environment, thereby improving reward prediction.
Figure 2:
4.2 Reward Redistribution based on Bidirectional Attention Mechanism
Building on Theorem 0 which establishes that optimal policies in the BRMDP and the MDP after reward redistribution equivalent, focus shifts to the crucial process of redistributing rewards within bag. To capture the contextual influence of each instance within the sequence, the Causal Transformer (Vaswani
ct al. 2017) naturally serves as suitable sequential prediction model. Traditionally, Transformers in RL have been used in a unidirectional manner (Chen el. al. 2021 Janner et. al. 2021 Micheli el al. 2023
where only previous instances influence the current prediction due to the unobservability of future instances. However, since our bagged rewards are generally non-Markovian, both procoding and subsoquent instances can influence the contribution of the current instance to the bagged reward. Therefore, understanding the relationships among instances within bag is therefore pivotal. This insight leads to employ bidirectional attention mechanism (Seo et al. 2016 Vaswani et al. 2017] Devlin et al. 2019) to capture these relationships. This mechanism connects both past and future instances within a bag, enabling a more comprehensive understanding of contextual influences. By quantitatively evaluating the contribution of each instance, the bidirectional attention mechanism facilitates nuanced reward redistribution, as demonstrated experimentally in Section
4.3 Reward Bag Transformer
We introduce the Reward Bag Transformer (RBT), novel reward redistribution method for the BRMDP. The RBT is engineered to comprehend the complex dynamics of the environment through bags and to preciscly predict instance-level rewards, facilitating effective reward redistribution.
4.3.1 Causal Transformer for Reward Representation
Referring to Fig. 2, the RBT comprises a Causal Transformer (Vaswani el. al. 2017 Radford et al. 2018 which maintains the chronological order state-action pairs (Chen et al.) 2021 Janner et al. 2021). For each time step
sequence of M time steps, the Causal Transformer, represented as function f. processes the input sequence {su, ag. S18-1, aM-1), generating the output )-M f(o). By aligning the output head with the action token at, we directly model the consequence of actions, which

---

Published in Transactions on Machine Learning Research (04/2025)
Algorithm Policy Optimization with RBT
Initialize replay buffer D, RBT parameters 8.
for trajectory collected from the environment do Store Irajectory with bag information {(Br,, R(Br,,) })} Bi,,, EB, in D.
EB,
4: Sample batches from D.
5: Estimate bag loss based on Eq. equation 7
6: Update RBT parameters based on the loss. Relabel rewards D using the updated RBT.
Optimize policy using the relabeled data by off-the-shelf RL algorithms (such SAC (Haarnoja al.
2018)). ).
end for
are pivotal computing immediate rewards and predicting subsequent states, thereby helping the model better understand environmental dynamics.
4.3.2 Bidirectional Attention Layer for Reward Redistribution
Once we have obtained the output embeddings (==) for reward prediction, they pass through bidi- rectional attention layer to produce (f where f; = (s( at) with 0 being the RBT parameters. This layer addresses the unidirectional limitation of the Causal Transformer architecture (Vaswani et al.) 2017; Radford et al. 2018). integrating past and future contexts for enhanced reward prediction accuracy. For state prediction, since the state at step t ++] depends only on the previous steps, is input into a state linear decoder, yielding the predicted next stale 8 = Sa(s,, 4z).
The core of the RBT architecture is its bidirectional attention mechanism. For each output embedding zt. we apply three different lincar transformations to obtain the query embodding qe E Rd key embedding k, E Rdd and value embedding vi R, where d is the embedding dimension of the key. The instance-level reward is then calculated by
M-1 N = Σ
{(g kp M )e Ug.
softmax
(4)
Vd
The rescaling operation is used to prevent extremely small gradients as in Vaswani et al. (011 This mechanism enables the RBT to consider both the immediate and contextual relevance of cach instance in a bag when predicting rewards.
4.4 Learning Objectives
The learning objectives of the RBT are twofold: (1) reward prediction within each reward bag and (2) state transition forecasling. These objectives are critical for enabling the model to navigate the complex dynamics of BRMDP environments.
4.4.1 Reward Prediction
The RBT is trained ensure that, for each reward bag, the sum of predicted instance-level rewards matches the total bagged reward. This is vital for maintaining the integrity of the reward structure in the BRMDP framework. The loss function is expressed
i Σ Pt
R(Bi,.e) )
L. Bi,n,)
(5)
This loss function encourages RBT to Iearn a meaningful distribution of rewards across instances within a bag while ensuring that the total redistributed reward remains consistent with the original bagged reward, By enforcing this consistency, the model aligns with the conditions established in Theorem

---

Published in Transactions on Machine Learning Research (04/2025)
4.4.2 State Transition Forecasting
Alongside reward prediction, the RBT is tasked with accurately predicting the next state in the environ- ment given the current stale and action. This capability crucial for understanding the dynamics of the environment, which in turn enhances the accuracy of reward prediction. The corresponding loss function is
C, St, de St+1) || s1+1ll?
(6)
where || - denotes the -nomm. This loss emphasizes the model's understanding of dynamics.
4.4.3 Composite Loss
The final learning objective combines the reward and state prediction losses:
Lbag = E [c, B Ban, E B.]
TーD
+ E Cx, Sc, ae, 81- (8t, Oes 8141) E ET], TND
where B > O is balancing cocfficient between the two loss components, and D denotes the replay buffer.
The RBT's s dual predictive capacity is its key advantage, enabling precise reward redistribution to individ- ual instances and forecasting the next state. This leverages environmental dynamics for enhanced reward distribution as experimentally evidenced in Section 5 Integrated with off-the-shelf RL algorithms such as SAC (Haarnoja et al. 2018). the RBT can enhance policy learning within the BRMDP framework, as outlined Algorithm
5 Experiment
In the following experiment section, we scrutinize the efficacy of our proposed method using benchmark tasks from both the MuJoCo (Brockman et al., 2016) and the DeepMind Control Suite (Tassa al.) 2018) environments, focusing scenarios with bagged rewards. Initially, we assess the performance of our method to understand its overall effectiveness. Subsequently, accurately predicts rewards. Finally, we evaluate the indispensability of cach component of the reward model, questioning if every part essential.
examine whether the proposed RBT reward model
5.1 Performance Comparison
To evaluate the performance of our method, we conduct a detailed comparison across multiple baselines. This analysis helps us understand how our approach fares against existing methods in handling scenarios with bagged rewards.
5.1.1 Experiment Setting
We evaluated our method on benchmark tasks from the MuJoCo locomotion suite (Ant-v2, Hopper-v2, HalfCheetah-v2, and Walker2d-v2) and the DeepMind Control Suite (cheetah-run, quadruped-walk, fish- upright, cartpole-swingup, ball ball_in__cm _i______uucccchh,, and reacher-hard).
In standard reinforcement learning settings, rewards are assigned at each step based on immediate state- action transitions. However, in our setup, only the bagged reward is observable, while the individual rewards for each instance remain unknown. Bagged rewards constructed by aggregating the rewards of all state- action pairs within a predefined sequence. Specifically, in the main text, present experimental results using coamation-bbsed aggregation, while additional results for more complex reward structures are provided in Appendix C To implement this reward structure, the bagged reward is assigned to the last instance within the bag, while all other state-action pairs within the bag receive a reward of zero. This setting effectively simulates delayed or sparse reward scenarios, making the learning problem more challenging due the lack of direct instance-level reward supervision. For each task, we collected 1e6 time steps, while the maximum episode length fixed at 1,000 steps, ensuring consistency across experiments.
8

---

Published in Transactions on Machine Learning Research (04/2025)
5.1.2 Baselines
In the comparative analysis, our framework was rigorously evaluated against several leading algorithms in the domain of RLTF:
SAC (Haarnoja et al.
2018) It directly utilized the original bagged reward information for policy training using the SAC algorithm.
IRCR Gangwani et al.
2020 It adopted a non non-parametric uniform reward redistribution ap- proach. We have adapted IRCR for bagged reward setting.
RRD [Ren et al.)
h021 It employed reward model trained with randomized return decompo- sition loss. We have adapted RRD for bagged reward setting.
HC (Han et al. 2022): The HC-decomposition framework was utilized to train the policy using a value function that operates on sequences of data. We employed the code as provided by the original paper.
Shaping (Hu et al., 2020 : A widely used reward shaping method for sparse reward setting, the
shaping reward function is adopted from Hu et al. (2020).
LIRPG (Zheng et al., 2018 It learned an intrinsic reward function to complement sparse envi- ronmental foedback, training policies to maximize combined extrinsic and intrinsic rewards. We use the same code provided by the paper.
While methods like RUDDER (Arjona-eedina et al. 2019) and Align-RUDDER (Patil et al. 2022) are known for addressing the problem of trajectory feedback, previous studies (Gangwani et al. 2020 Ren et al. Zhang et al. 2023) have shown superior performance using referenced methods. Additionally, since Align-RUDDER RÜDDER relics on successful trajectorics for scoring state-action pairs, which is impractical in MuJoCo
2021
Patil et al. 2022), we ultimately excluded both methods from our comparison. Sparse reward baselines were included in our comparisons to highlight the distinctions and challenges unique to our bagged reward setting, While both settings provide rewards after sequence, sparse rewards are Markovian, whereas our bagged rewards are non-Markovian. Therefore, including sparse reward baselines illustrates that directly applying sparse reward methods to bagged rewards is possible but not effective. For more discussion sparse rewards, see Appendix B Besides, detailed descriptions of the model parameters and hyper-parameter used during training are provided in Appendix A with more experimental results available in Appendix C
5.1.3 Evaluation Metric
We report the average accumulative reward across seeds with random initialization to demonstrate the performance of evaluated methods. Higher accumulative reward in evaluation indicates better performance.
5.1.4 Experimental Results
In the fixed-length reward bag experiment (see Fig. 3 25, 50, 100, 200, and 500) and trajectory feedback (labeled as 9999) across each environment, where bagged reward made by accumulating immodiate rewards within bags. This aimed to illustrate the influence of varying bag lengths, providing insight into how bag size affected the performance of the learning algorithm. The SAC method, using bagged rewards directly from the environment, suffers from lack of guidance in agent training due to missing immediate reward information. This issue worsens with longer bag lengths, indicating that increased reward sparsity leads to less effective policy optimization. The IRCR and RRD methods, treating rewards uniformly within a reward bag, oulperform SAC, suggesting benefits from even approximate guidance. However, notable variance their results indicates potential consistency and re- liability issues. The HC method excels only with shorter bag lengths, suggesting that this value function modification method struggles utilize information from longer sequences. The Shaping and LIRPG meth- ods exhibit. subdued performance across tasks, as they are designed to solve sparse reward problems with
we conducted experiments with six bag lengths (5,
9

---

Published in Transactions on Machine Learning Research (04/2025)
Hopper-v2
HellCheetah--2
Walker2d
01.01
10:01
10:03
-004
200
60002
Netur
1700
2000
-2233
103
300
bag Langth
Bag Lengon
Bag Langth
Bag Length
cheetah-sun
quadruped-walkk
fich-upright
Carppole-eenguup
8ebu
6:0
300
180
202
120
100
123
Figure 3: Performance comparison across fixed-length reward bag in MuJoCo (top row) and DeepMind Control Suite (bottom row) environments with six different length bag settings (5, 25, 50, 100, 200, and 500) and trajectory feedback (labeled as 9999). The mean and standard deviation are computed over 6 trials with different random seeds across a total of 1e6 time steps. Additionally, the black dashed line the figures represents the performance of vanilla SAC with a bag length of along with its standard deviation, serving reference bascline.
Markovian rewards, which do not align with the reward bag setting. The proposed RBT method consistently outperforms the other approaches across all the environments and bag lengths, showing that it is not only well-suited for environments with short reward bags but also capable of handling large reward bag scenarios. This demonstrated the capability of RBT to learn from the sequence of instances and, by integrating bagged reward information, accurately allocate rewards to instances, thereby guiding better policy training. No- tably, for shorter bag lengths, the performance of the RBT method elosely matches that of vanilla SAC. This suggests that RBT effectively preserves critical reward signals, allowing the agent to achieve performance comparable to the vanilla SAC baseline. In addition to experiments conducted under settings consistent with previous work (Gangwani et al. 2020 Ren et al., 2021 Zhang et al. 2023),, where the bagged reward is the sum of immediate rewards, we also tested bagged rewards with more complex structures (sec Appendix C.2 These results further demonstrated that our method out performs the baselines, indicating that previous approaches are effective only when bagged rewards are the sum of immediate rewards, highlighting their limitations.
To validate the effectiveness of our approach under more complex conditions, we designed experiment with arbitrary-length reward bags, allowing not only for variations in bag lengths but also for overlaps and gaps between reward bags. This setup simulated more realistic scenarios and tested the robustness of our method. The results, detailed in Table D confirm the superior performance of the proposed RET method in these complex reward settings. These findings suggest RBT's potential for broader applications, demonstrating its versatility and robustness in handling intricate reward dynamics.
We also conducted experiments in sparse reward environments, where tasks are rewarded with 1 at the end of the trajectory if completed and 0 otherwise. In this setting, the reward does not accumulate in redistribution model. As shown in Table 2 RLTF methods perform poorly in sparse reward environments. In contrast, our proposed method not only outperforms RLTF baselines but also exceeds the performance of sparse reward baselines. From another perspective, we also tested the performance of sparse reward methods in the bagged reward sctting (sce Appendix B The results indicated that, due to differences the reward structure, sparse reward methods do not adapt well in the bagged reward setting either. This demonstrates that our proposed RBT model can effectively capture relationships between instances across different reward
10

---

Published in Transactions on Machine Learning Research (04/2025)
Table 1: Performance comparison across arbitrary-lentth reward bag confgurattooss over trials with 1e6 time steps for training, presenting average scores and standard deviations. "Narrow" refers to bags with lengths varying arbitrarily from 25 to 200 and narrow intervals between -10 to 10. "Broad" denotes the setting with bag lengths varying arbitrarily from 100 to 200 and broad interval variations from -40 to 40. The best and comparable methods based on the paired t-test at the significance level 5% are highlighted in boldface.
Bag Setting
Environment
SAC
IRCR
RRD
IIC
Shaping
LIRPG
(oT((
0.87
368.69
2272.39
106.92
-732.03
-756.78
5122.50
Ant-v2
(2.98)
(119.74)
(835.86)
(153.86)
(225.75)
(763.66)
(206.44)
Hopper-v2
317.72
3353.35
2184.41
510.66
1031.38
126.13
3499.54
Narrow
(52.17)
(61.97)
(807.71)
(94.49)
(335.05)
(30.18)
(76.62)
788.45
10853.85
9709.62
4027.25
2348.78
1101.38
11282.24
HalfCheetab-v2
(1737.57)
(573.72)
(1479.73)
(441.01)
(390.48)
(1248.45)
(266.08)
193.07
4144.65
3536 .90
309.19
1022.32
123.43
4983.39
Walker2d-v2
(48.40)
(673.66)
(546.66)
(171.69)
(369.81)
(50.97)
(311.09)
-3.31
368.69
1323.50
5.97
-634.69
1264.08
5167.79
Ant-v2
(1.15)
(158.46)
(1079.60)
(20.08)
(161.57)
(416.86)
(303.83)
329.48
3296.20
1102.38
701.84
1066.51
203.01
3499.53
Hopper-v2
Broad
(44.21)
(216.35)
(892.12)
(149.44)
(429.68)
(177.80)
(94.00)
HalfChoetah-v2
43.96
9158.14
4199.16
4460.80
1459.83
924.26
10837.15
(94.32)
(1402.62)
(1476.85)
(518.94)
(1251.52)
(1110.97)
(254.99)
176.09
4179.08
330.96
447.45
947.87
194.95
5202.38
Walker2d-v2
(49.81)
(937.42)
(79.26)
(155.63)
(90.46)
(98.05)
(248.35)
Table 2: Performance comparison sparse reward enviromments over 6 trials with 1e6 time steps for training, presenting average scores and standard deviations. Rewards are given with at the end of the trajectory if completed and 0 otherwise. The best and comparable methods based on the paired t-test at the significance level 5% are highlighted in boldface.
Environment
SAC
IRCR
RRD
HC
Shaping
LIRPG
RBT(ours)
ball cup-catch
646.90
773.19
666.18
76.19
847.88
125.50
972.19
(231.15)
(73.06)
(33.02)
(67.25)
(178.01)
(21.25)
(5.50)
531.88
584.76
347. 16
7.57
664.49
8.74
735.67
peacher-aard
(127.94)
(269.32)
(106.56)
(3.50)
(170.50)
(2.36)
(144.90)
structures. By accurately distributing reward signals across instances, it provides more informative feedback, enabling the agent to learn more efficiently and ultimately improving overall policy performance.
5.2 Case Study
The previous experimental results showcase the superiority of RBT over baselines. This led to intriguing inquiry: Is the RBT reward model proficient in accurately redistributing rewards? To investigate this question, we performed an experiment focused on reward comparison, utilizing a trajectory generated by an agent trained in the Hopper-v2 environment with a bag length of 100, where the bagged reward represents an accumulated reward. As shown Fig. 4. which spans 1000 steps, RBT-predicted rewards, unobservable true rewards, and observable bagged rewards (presented in uniform format for better visualization) are compared. The figure indicates that the rewards predicted by the RBT closely match the trends of the true rewards. This observation suggest that the RBT is effective at reconstructing true rewards from bagged rewards, despite the coarse nature of the environmental reward signals.
Beneath the figure, a series of images depicts a complete jump cycle by the agent, illustrating its motion sequence: mid-air, landing, jumping, and returning to mid-air. Red boxes highlight specific states that correspond to reward peaks and troughs, representing moments of maximum, minimum, and moderate rewards. In the Hopper-v2 environment, rewards consist of constant "healthy reward" for operational integrity, a "forward reward" for progress in the positive x-direction, and "control cost" for penalizing large
11

---

Published in Transactions on Machine Learning Research (04/2025)
Rewards Comparison
Reward
3
Pred Reward
True Reward (unobservable)
Bagged Reward (observable)
200
400
600
800
1000
Step
220
235
265
Figure Rewards comparison and agent states a trajectory with a bag length of 100 in the Hopper- environment. The top graph compares predicted rewards against true rewards and aggregated bagged rewards. The bottom images show the Hopper agent's motion at different time steps, with the highlighted frames representing key moments that provide insight into the relationship between agent behavior and reward dynamics.
Ant-v2
Hopper-u2
HallCheetahh-22
Walker2d-42
4035
10000
5.000
0208
102
431
c00
030
Langth
Bag
Langth
Langth
Figure 5: Ablation study of reward model components across various environments. The chart presents mean and standard deviation of rewards over trials with 1e6 timesteps, showcasing the efficacy of the full proposed method relative to its variants without certain features.
actions. At peak reward instances, the agent is typically fully grounded an optimal posture for forward leaping, which maximizes the "forward reward" through pronounced x-direction movement. Concurrently, it sustains the "healthy reward" and minimizes "control cost" through measured, efficient actions. This analysis underscores that the RBT can adeptly decode the environmental dynamics and the muanced reward redistribution even under the setting of RLBR.
5.3 Ablation Study
We conducted comparisons to examine the role of RBT's modules. As shown in Fig. the full RBT model consistently outperforms its variants, indicating synergistic effect when all components together. Performance drops significantly when the bidirectional attention mechanism is removed, especially in complex environments like Ant-v2 and HalfChectah-v2, suggesting its critical role in accurate reward prediction. Additionally, we can observe that removing the next state prediction component weakens the reward model's understanding of environmental dynamics, reducing reward prediction accuracy and hin- dering policy learning. The greatest performance decline occurs when both the next state prediction and bidirectional self-attention mechanism are absent, underscoring their individual and combined importance in building robust reward model.
used
12

---

Published in Transactions on Machine Learning Research (04/2025)
6 Discussion
In this section, we discuss the limitations of our approach, well as its broader applicability and potential directions for future research.
6.1 Limitations
One major limitation of our approach is the training duration, which is longer compared to methods that do not require a reward model (see Appendix C.5 The increased computational cost primarily arises from the reward relabeling process, where the model updates assigned rewards in the replay buffer based on the lcarned reward distribution. This step adds siguificant overhead, making training less efficient, particularly for longer bag lengths, where reward prediction becomes more challenging. A potential direction for improving computational efficiency is prioritized experience replay (Schaul et al.] 2015), which could selectively relabel more informative samples rather than processing the entire buffer equally, thereby reducing unnecessary computation. Another possible enhancement is cached reward update mechanism, where only a subset of stored samples is updated per iteration instead of recomputing rewards for the full replay buffer. These optimizations have the potential to reduce training time while maintaining learning performance.
Another limitation of our approach is the assumption that the bag partition function 9 provided by the environment. While this assumption simplifies the problem setting and allows us focus on reward redistribution, it may not always hold in real-world scenarios. In cases where bag boundaries are unknown, additional inference step would be required to identify bag transitions, potentially increasing the complexity of both learning and computational cost. Future work could explore methods for learning bag partitioning dynamically, such as incorporating segmentation models (Adams & MacKay, 20077] Truong et al. 2020),
leveraging latent variable approaches (Yu 2010 Kingma et al. 2013 Johnson Willsky 2013 or attention- based mechanisms (Yi et al. 2021) to infer bag structures from data.
Additionally, our current study does not explicitly address the generalization ability of the learned reward redistribution model across different environments. All experiments involve training and evaluating policies within individual environments without testing domain transfer capabilities. While this setup ensures controlled evaluation, it remains unclear how well the model would adapt to new domains with different reward structures. Future work could investigate techniques such as meta-learning (Finn el al. 2017 Zintgraf
et al. 2019 or domain adaptation (Higgins et al. 2017 Peng et al., 2018 to improve generalization.
Finally, our experiments primarily focus on the Control Suite. While this setting provides a clean testbed for evaluating our method, real-world problems often involve partial observability, sharp reward transitions, and challenging exploration scenarios. In environments where observations are incomplete or noisy, additional
techniques such as belief state estimation (Kaelbling et al. 1998 Murphy 2000 [g] et al.. 2018 may be required to improve robustness. For environments with sparse and sharp reward transitions, a possible approach is to define bag boundaries based on reward occurrences, as done in our sparse reward experiments.
Yu 2010 Kingma et al., 2013 Johnson & Willsky 2013) could be explored to better align reward redistribution with critical decision points. Lastly, our method focuses on credit assignment rather than exploration. In environments with
Morever, adaptive bag partitioning (Adams MacKay 2007 Truong et al. 2020
difficult exploration, additional techniques such as intrinsic motivation (Schmidhuber] 1991 Pathak et al., 2017
Burda. et. al.) 2019a) or entropy-based exploration (Hazan el. al. 2019 Lee el al. 2021 could be
integrated to improve performance.
6.2 Broader Applicability and Future Directions
A promising direction for extending our work is to apply the RLBR framework to practical scenarios involving both embodied and digital agents. Many real-world tasks in these domains can be naturally segmented into contiguous phases, each consisting of sequence of interrelated actions that jointly contribute to intermediate progress toward the final goal (Kulkarni et al.) 2016 Nachum et al., 2018). In embodied agents (Parr & Russell 1997), such as houschold robots, tasks like cleaning, assembling, or navigation consist of successive physical interactions that can be grouped into temporally coherent segments based on spatial or behavioral continuity, making them well-suited for modeling with the RLBR framework. In digital agents (Peng et al.
13

---

Published in Transactions on Machine Learning Research (04/2025)
2017 Shi et al. 2019 such dialogue systems web-based task solvers, feedback is often delayed and context-dependent, typically emerging only after sequence of actions achieves sub-goal. For example, positive feedback may appear only after several dialogue turns. This makes per-action credit assignment challenging, whereas attributing rewards to contiguous segments better captures their joint contribution and improves learning stability. This structure aligns well with the bagged reward lormulation, where reward signals are assigned to temporally localized segments rather than individual steps or the entire trajectory.
intermediate form of reward between instance-level and trajectory-level feedback, holds strong potential for real-world applica- tions. It can serve bridge between low-level control and high-level task abstraction. Future work may explore integrating RLBR with hierarchical reinforcement learning (Barto & Mahadevan, 2003 Nachum jet al. 2018) or jointly learning bag partitioning and policy optimization to support more adaptive behavior in complex domains.
Given this structural alignment, believe the RLBR framework, which focuses
7 Conclusion
this paper, we introduced a novel learning framework called Reinforcement Learning from Bagged Rewards (RLBR). To address the challenges posed by this framework, we established theoretical connections bet ween Bagged Reward MDPs (BRMDPs) and original MDPs after reward redistribution, providing a solid foun- dation for approach. Based on this theory, we proposed reward model, the Reward Bag Transformer (RBT), designed to redistribute rewards by interpreting contextual ual information within bags and understand- ing environmental dynamics. The efficacy of RBT was demonstrated through extensive experiments, where it consistently outperformed existing methods across various reward bag scenarios. Additionally, our case studics highlighted RBT's ability to effectively reallocate rewards while maintaining fidelity to the underlying true reward structure.
Acknowledgments
YT was supported by Institute for Al and Beyond, UTokyo. XC was supported by JSPS, KAKENHI Grant Number JP24KJ0610, Japan. YD was supported by National Natural Science Foundation of China, Grant Number 62206245. MS was supported by JST ASPIRE Grant Number JPMJAP2405, and Institute for Al and Beyond, UTokyo. The authors would like to thank anonymous reviewers for their insightful reviews.
14

---

Published in Transactions on Machine Learning Research (04/2025)
References
Ryan Prescott Adams and David JC MacKay, Bayesian online changepoint detection. ar preprint
arXiv:0110.9742, 2007.
Pranav Agarwal, Pierre de Beaucorps, and Raoul de Charette. Goal-constrained sparse reinforcement learn- ing for end-to-end driving. arXie preprinl arXiv:2103.09180 2021.
Jose A Arjona-Medina, Michael Gillhofer, Michael Widrich, Thomas Unterthiner, Johannes Brandstetter, and Sepp Hochreiter. Rudder: Return decomposition for delayed rewards. In The Thirty-second Annual Conference on Advances in Newral Information Processing Systems, 2019.
Andrew G Barto and Sridhar Mahadevan. Recent advances in hierarchical reinforcement learning. Discrete event dynamic systems, 13:341 379, 2003.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym, 2016.
Yuri Burda, Harri Edwards, Deepak Pathak, Amos Storkey, Trevor Darrell, and Alexei A Efros. Large-scale study of curiosity-driven learning. In International Conference on Learning Representations, 2019a.
Yuri Burda, Harrison Edwards, Amos Storkey, and Oleg Klimov. Exploration by random network distillation. In The Seventh International Conference on Learning Representations, 2019b.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. In The Thirty-fourth Annual Conference on Advances in Neural Information Processing Systems, 2021.
Sirui Chen, Zhaowei Zhang, Yaodong Yang, and Yali Du Stas: Spatial-temporal return dccomposition for solving sparse rewards problems in multi-agent reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pp. 17337-17345, 2024.
Jacob Devlin, Ming-Way Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidi- rectional transformers for language understanding. In Jill Burstein, Christy Doran, and Thamar Solorio (ods.). Proceedings of the 2019 Conference of the North American Chapter of the Association for Com- pulational Lánguisties: Human Language Technologies, NAACL-HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume 1 (Long and Short Papers), pp. 1771-116. Association for Computational Linguistics, 2019.
Joseph Early, Tom Bewley, Christine Evers, and Sarvapali Ramchurn. Non-markovian reward modelling from trajectory labels via interpretable multiple instance learning. In The Thirby-fifth Annual Conference on Advances in Neural Informalion Processing Systems, 2022.
Chelsea Finn, Picter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. Tn International conference on machine learning, pp. 1126-1135. PMT.R, 2017.
Carlos Florensa, David Held, Xinyang Geng, and Pieter Abbeel. Automatic goal generation for reinforcement learning agents. In The Thirty-fifth International Conference on Machine Learning. PMLR, 2018.
Roy Frostig, Matthew James Johnson, and Chris Leary. Compiling machine learning programs via high-level tracing. Systems for Machine Learning, 4(9), 2018.
Tanmay Gangwani, Yuan Zhou, and Jian Peng. Learning guidance rewards with trajectory-space smoothing. The Thirty-third Anreual Conference on Advances in Neural Information Processing Systems, 2020.
Daiqi Gao, Hsin-Yu Yu Lai, Predrag Klasnja, and Susan A Murphy. Harnessing causality in reinforcement learning with bagged decision times. arXiv preprint (Xii:2210011652 2024.
Maor Gaon and Ronen Brafman. Reinforcement learning with mon-vvvoovian rewards. In The Chirty-Fowrde
AAAT Conference on Artificial Intelligence, volume 34, 2020.
15

---

Published in Transactions on Machine Learning Research (04/2025)
Tuomas Haarnoja, Aurick Zhou, Picter Abbeel, and Sergey Levine. Soft actor-critic: Oft-policy maximum entropy deep reinforcement learning with stochastic actor. In The Thirly-fifhh International Conference on Machine Learning. PMLR, 2018.
Beining Han, Zhizhou Ren, Zuofan Wu, Yuan Zhou, and Jian Peng. Oft-policy reinforcement learning with delayed rewards. In The nhith International Conference on Machine Learning. PMLR, 2022.
Elad Hazan, Sham Kakade, Karan Singh, and Abby Van Soest. Provably efficient maximum entropy explo-
ration. In International Conference on Machine Learning, pp. 2681 2691. PMLR, 2019.
Irina Higgins, Arka Pal, Andrei Rusu, Loic Matthey, Christopher Burgess, Alexander Pritzel, Matthew Botvinick, Charles Blundell, and Alexander Lerchner. Darla: Improving zero-shot transfer in reinforcement learning. In International conference on machine learning, pp. 1480- 1490. PMLR, 2017.
Yujing Hu, Weixun Wang, Hanglian Jia, Yixiang Wang, Yingleng Chen, Jianye Hao, Feng Wu, and Changjie Fan. Learning to utilize shaping rewards: A new approach of reward shaping. In The Thirty- thitt- Annual Conference on Advances in Newral Information Processing Systems, 2020.
Maximilian Igl, Luisa Zintgraf, Tuan Anh Le, Frank Wood, and Shimon Whiteson. Deep variational rein- forcement lcarning for pomdps. In International conference on machine learning, pp. 2117-2126. PMLR, 2018.
Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. In The foirth-fourth Annual Conference on Advances in Neural Information Processing Systems, 2021.
Matthew J Johnson and Alan S Willsky. Bayesian nonparametric hidden semi-markov models. The Journal of Machine Learning Research, 14(1):673 701, 2013.
Leslie Pack Kaelbling, Michael I. Littman, and Anthony R Cassandra. Planning and acting in partially observable stochastic domains. Artificial intelligence, 101(1-2):99 134, 1998.
Nan Rosemary Ke, Anirudh Goyal ALIAS PARTH GOYAL, Olexa Bilaniuk, Jonathan Binas, Michael C Mozer, Chris Pal, and Yoshua Bengio. Sparse attentive backtracking: Temporal credit assignment through reminding. In The Thirty-first Annual Conference Advances in Neural Information Processing Systems, 2018.
Diederik P Kingma, Max Welling, et al. Auto-encoding variational bayes, 2013.
B Ravi Kiran, Ibrahim Sobh, Victor Talpaert, Patrick Mannion, Ahmad A Al Sallab, Senthil Yogamani, and Patrick Pérez. Deep reinforcement learning for autonomous driving: A survey. IEEE Transactions on Intelligent Transportation Systems, 23(6):4909 4926, 2021.
llya Kostrikov. JAXRL: Implementations of Reinforcement Learning algorithms in JAX, 10 2021. URL
Tejas D Kulkarni, Karthik Narasimhan, Ardavan Saeedi, and Josh Tenenbaum. Hierarchical deep reinforce- ment learning: Integrating temporal abstraction and intrinsic motivation. Advances in neural information processing systems, 29, 2016.
Minae Kwon, Sang Michael Xie, Kalesha. Bullard, and Dorsa Sadigh. Reward design with language models. In The Eleventh International Conference on Learning Representations, 2023.
Harrison Lee, Samrat Phatale, Hassan Mansoor, Kellie Lu, Thomas Mesnard, Colton Bishop, Victor Car- bune, and Abhinav Rastogi. Rlaif: Scaling reinforcement learning from human feedback with ai feedback, 2023.
Kimin Lee, Michael Laskin, Aravind Srinivas, and Picter Abboel. Sunriso: A simple unified framework for ensemble learning in deep reinforcement learning. In International Conference on Machine Learning, pp. 6131-6141. PMLR, 2021.
16

---

Published in Transactions on Machine Learning Research (04/2025)
Yang Liu, Yunan Luo, Yuanyi Zhong, Xi Chen, Qiang Liu, and Jian Peng. Sequence modeling of temporal credit assignment for episodic reinforcement learning, 2019.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In The Stath International Conference on Learning Representations, 2018.
Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are sample-cfficient world models. In The Eleventh International Conference on Learning Representations, 2023.
Kevin P Murphy. A survey of pomdp solution techniques. environment, 2(10). 2000.
Ofrr Nachum, Shixiang Shane Gu, Honglak Lee, and Sergey Levine. Data-efficient hicrarchical reinforcement learning. Advances in newral information processing systems, 31, 2018.
Andrew Y Ng. Daishi Harada, and Stuart Russell. Policy invariance under reward transformations: Theory and application to reward shaping. In The Sirteenth International Conference on Machine Learning. PMLR, 1999.
Ronald Parr and Stuart Russell. Reinforcement learning with hierarchies of machines. Advances in newral information processing systems, 10, 1997.
Deepak Pathak, Pulkit Agrawal, Alexei A Efros, and Trevor Darrell. Suriosity-driver exploration by self- supervised prediction. In The Thirty-fourth International Conference on Machine Learning. PMTR, 2017.
Vihang Patil, Markus Hofmarcher, Marius-Coon Constantin Dinu, Matthias Dorfer, Patrick M Blies, Johannes Brandstetter, José Arjona-Meding and Sepp Hochreiter. Align-rudder: Learning from few demonstrations by reward redistribution. In The Thirty-ninth International Conference on Machine Learning. PMLR, 2022.
Baolin Peng, Xiujun Li, Lihong Li, Jianfeng Gao, Asli Celikyilmaz, Sungjin Lee, and Kam-Fai Wong. Composite task-completion dialogue policy learning via hicrarchical deep reinforcement learning. arXiv preprint Xrxi::1704.00:0 080084 2017.
Xue Bin Peng, Marcin Andrychowicz, Wojciech Zaremba, and Pieter Abbeel. Sim-to-real transfer of robotic control with dynamics randomization. In 2018 IEEE international conference on robotics and automation (ICRA), pp. 3803-3810. IEFE, 2018.
Alec Radford, Karthik Narasimhan, Tim Salimans, Tya Sutskever, et al. Improving language understanding by generative pre-training, 2018.
Zhizhou Ren, Ruihan Guo, Yuan Zhou, and Jian Peng. Learning long-term reward redistribution via ran- domized return decompasition. In The Ninth International Conference on Learning Representations, 2021.
Martin Riedmiller, Roland Hafner, Thomas Lampe, Michael Neunert, Jonas Degrave, Tom Wiele, Vlad Mnih, Nicolas Heess, and Jost Tobias Springenberg. Learning by playing solving sparse reward tasks from scratch. In The Thirty-fifth International Conference on Machine Learning. PMLR, 2018.
Tom Schaul, John Quan, Ioannis Antonoglou, and David Silver. Prioritized experience replay. arXiv preprint arXiv:1511.0599 2015.
Jürgen Schmidluber Curious model-buildigg control systems. In Proc. international joint conference on neural networks, pp. 1458 1463, 1991.
Ramanan Sekar, Oleh Rybkin, Kostas Daniiidiis, Pieter Abbeel, Danijar Hafner, and Deepak Pathak. Plan- ning to explore via self-supervised world models. In The shirth-seventh International Conference on Machine Learning. PMLR, 2020.
Minjoon Sco, Aniruddha Kembhavi, Ali Farhadi, and Hannanch Hajishirzi. Bidirectional attention flow for machine comprehension. In The Fourth International Conference on Learning Representations, 2016.
17

---

Published in Transactions on Machine Learning Research (04/2025)
Weiyan Shi, Tiancheng Zhao, and Zhou Yu. Unsupervised dialog structure learning. arXiv preprint arXiv:1904 003736 2019.
David Silver, Aja Huang, Chris Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, Toannis Antonoglou, Veda Panneershelvam,, Mare Lanctot, et al. Mastering the game of go with deep neural networks and tree scarch. nature, 529(7587):484-489 2016.
Jonathan Sorg, Richard L Lewis, and Satinder Singh. Reward design via online gradient ascent. In The Twenty-fourth Annual Conference on Advances in Neural Information Processing Systems, 2010.
R.S. Sutton and A.G. Barto. Reinforcement learning: An introduction. IEEE Transactions on Neural Networks, 9(5):1054 -054,, 1998.
Pradyumna Tambwekar, Murtaza Dhuliawala, Lara J. Martin, Animesh Mchta, Brent Harrison, and Mark O. Riedl. Controllable neural story plot generation via reward shaping. In Proceedings of the International Joint Conference on Artificial Intelligence. International Joint Conferences on Artificial Intelligence Organization, 2019.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq. et al. Deepmind control suite, 2018.
Chen Tessler, Daniel J Mankowitz, and Shie Mannor. Reward constrained policy optimization. In The Seventh International Conference on Learning Representations, 2019.
Charles Truong, Laurent Oudre, and Nicolas Vayatis. Selective review of offiin change point detection methods. Signal Processing, 167:107299, 2020.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In The Thirtieth Annual Conference on Advonces in Neural Information Processing Systems, 2017.
Oriol Vinyals, Igor Babuschkin, Wojeiech M Czarnecki, Michaël Mathieu, Andrew Dudzik, Junyoung Chung, David H Choi, Richard Powell, Timo Ewalds, Petko Georgiev, et al. Grandmaster level in starcraft ii using multi-agent reinforcement learning. Nature, 2019.
Jingkang Wang, Yang Liu, and Bo Li. Reinforcement learning with perturbed rewards. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pp. 6202 6209, 2020.
Christopher John Cornish Hellaby Watkins, Learning from delayed rewards. King's College, Cambridge United Kingdom, 1989.
Michael Widrich, Markus Hofmarcher, Vihang Prakash Patil, Angela Bitto-Nemling, and Sepp Hochreiter, Modern hopfield networks for return decomposition for delayed rewards. In Deep RL Workshop NewrIPS 2021, 2021.
Peter Wurman, Samuel Barrett, Kenta Kawamoto, James MacGlashan, Kaushik Subramanian, Thomas J Walsh, Roberto Capobianco, Alisa Devlic, Franziska Eckert, Florian Fuchs, et al. Outracing champion gran turismo drivers with deep reinforcement learning. Nature, (02(7896):223-222) 2022.
Hongyang Yang, Xiao-Yang Liu, Shan Zhong, and Anwar Walid. Deep reinforcement learning for automated stock trading: An ensemble strategy. In Proceedings of the first ACM international conference on Al in finance, 2020.
Fangqiu Yi, Hongyu Wen, and Tingting Jiang. preprint ar Xiv:21 10. 08568, 2021.
Asformer: Transformer for action segmentation. arXiv
Chao Yu, Jiming Liu, Shamim Nemati, and Guosheng Yin. Reinforcement learning in healthcare: A survey. ACM Computing Surveys (CSUR), 55(1):1-3 (11:1-36,,,) 2021.
Shun-Zheng Yu. Hidden semi-markov models. Artificial intelligence, 174(2):215 243, 2010.
18

---

Published in Transactions on Machine Learning Research (04/2025)
Yudi Zhang, Yali Du, Biwei Huang, Ziyan Wang, Jun Wang, Meng Fang, and Mykola Pechenizkiy. Inter- pretable reward redistribution in reinforcement learning: A causal approach. In The Thirby-siath Annual Conference on Advances in Neural Information Processing Systems, 2023.
Lulu Zheng, Jiarui Chen, Jianhao Wang, Jiamin He, Yujing Hu, Yingfeng Chen, Changjie Fan, Yang Gao, and Chongjie Zhang. Episodic multi-agent reinforcement learning with curiosity-driven exploration. In The Thirly-fourth Annual Conference on Advances in Newral Information Processing Systems, 2021.
Zeyu Zheng, Junhyuk Oh, and Satinder Singh. On learning intrinsic rewards for policy gradient methods. In The Thirly-frrs Ansval Conference on Advances in Newral Information Processing Systems, 2018.
Luisa Zintgraf, Kyriacos Shiarli, Vitaly Kurin, Katja Hofmann, and Shimon Whileson. Fast context adap- tation via meta-leant leatt-lerrning.. In International conference on machine learning, pp. 7693-777777 7702. PMLR, 2019.
19

---

Published in Transactions on Machine Learning Research (04/2025)
A Experiment Settings and Implementation Details
Benchmarks with Bagged Rewards. We introduced a novel problem setting in the suite of MuJoCo and DeepMind Control Suite locomotion benchmark tasks, termed as bagged rewards. Our simulations ran on the OpenAT Gym platform (Brockman et al. 2016) and the DeepMind Control Suite (Tssa et al., 2018).
featuring tasks that stretched over long horizons with a set maximum trajectory length of T 1000. Wo used MuJoCo version 2.0 for our simulations, which is available at licensed under a commercial license, and we have adhered to its terms of service and licensing agreements stated on the official website. The DeepMind Control Suite is available under Apache License 2.0, and have complied with its terms of use.
org/ MuJoCo is
Reward bag experiments of different bag sizes (5, 25, 50, 100, 200, and 500) and trajectory feedback were set up to verify the effectiveness of the method. To evaluate the efficacy of proposed method, commonly used trajectory feedback algorithms were adapted to fit the bagged reward setting as baselines. In these experiments, each reward bag was treated as an individual trajectory, and these modified algorithms were applied accordingly. Additionally, experiments using standard trajectory feedback were conducted provide comparative bascline within the unique setting. The total episodic feedback was computed at the end of the trajectory and was the sum of the per-step rewards the agent had collected throughout the episode. This experiment setting was the same as some previous works for learning from trajectory feedback (Gangwani et al. 2020 Ren et al. 2021h.
For the sparse reward experiments, the ball in cup-catch task provides a sparse reward of 1 when the ball is the cup and otherwise; the reacher-hard task, the reward is 1 when the end effector penetrates the target sphere (Tassa al. 2018). Therefore, we treated the sparse reward as trajectory feedback and conducted experiments where the entire trajectory was considered as a single reward bag in these environments. The bagged reward was identical to the sparse reward provided by the environment.
Table 3: Hyper-parameter of RBT.
Hyper-parameter
Value
Number of Causal Transformer layers
Number of bidirectional attention layers
Number of attention heads
Embedding dimension
256
Batch size
64
Dropout rate
0.1
Learning rate
0.0001
Optimizer
AdamW Loshchilov & Hutter 2018
Weight decay
0.0001
Warmup steps
100
Total gradient steps
10000
Implementation Details and Hyper-parameter Configuration. In our experiments, the policy opti- mization module was implemented based on soft actor-critic (SAC) (Haarnoja et al. 2018). We evaluated the performance of our proposed methods with the same configuration hyper-parameters all environments. The back-end SAC followed the JaxRL implementation [Kostrikov 2021 which is available under the MIT License.
The RBT reward model was developed based on the GPT implementation in JAX Frostig et al. 2018) which is available under the Apache License 2.0. Our experiments utilized the Causal Transformer with three layers and four self-attention heads, followed by a bidirectional sell-attention layer with one sell-attention head. For detailed
settings of the RBT, please refer to Table
For the baseline methods, the IRCR (Gangwani al. 2020) method was implemented based on the de- scriptions provided in the original paper. The RRD (Ren al. 2021 and LIRPG (Zheng et. al. 2018
20

---

Published in Transactions on Machine Learning Research (04/2025)
Ant-v2
Hopper-v2
HallCheetah-22
Walker2div2
0.0
253
-2000
1090
-520
20
503
150
3020
100
999
Figure 6: Comparison of proposed method and sparse reward baselines. The mean and standard deviation are displayed over 6 trials with different random seeds across total of 1e6 time steps.
methods are both licensed under the MIT Licenso. The code of HC HHan al. 2022 is available in the supplementary material at htttps://openrrviie net:forrmmsddns
To ensure uniformity in the policy optimization process across all methodologies, each was subjected to 1, 1,000,000 training iterations. For the proposed method, we initially collated a dataset comprising 10,000 time steps to pre-train the reward model. This model then underwent 100 pre-training iterations, a step deemed essential adequately prepare the reward model before embarking on the principal policy learning phase. Following this initial warm-up period, the reward model trained for 10 iterations after each new trajectory incorporated. Morcover, systematically gauge performance and progress, evaluations were carried out at intervals of every 5,000 time steps. The computational resources for these procedures were NVIDIA GeForce RTX 2080 Ti GPU clusters with 8GB of memory, dedicated to training and evaluating tasks.
B Discussion on Sparse Reward
The sparse reward setting presents significant challenges due to infrequent feedback, making it difficult for agents to effectively explore the environment. and discover successful strategies. To address this challenge, various methods have been developed to enhance exploration. Reward shaping strategies (Ng et al.
1999 Hu
ct al. 2020 Tambwekar al. 2019) added rewards to actions in a way that guides the agent towards better
policies without altering the original reward function. Curiosity-driven methods (Pathak et al. 2017 Sekar ct al. 2020) encouraged agents to explore the environment by visiting unscen states, potentially solving
tasks with sparse rewards. Additionally, curriculun learning in RL (Florensa et al. 2018 Riedmiller et al. 2018) involved presenting an agent with a sequence of tasks with gradually increasing complexity, allowing the agent to eventually solve the initially given sparse reward task.
Due to the formal similarity between sparse rewards and trajectory feedback, both involving a single reward signal after a sequence, these methods can be directly applied to RLTF. However, the key difference lies the nature of the rewards: sparse rewards typically focus on specific key points, making the reward Markovian, whereas RLTF evaluates the entire trajectory, indicating a Non-Markovian reward. Furthermore, sparse reward strategies address the exploration-exploiaalee focuses on distributing the evaluation of the entire trajectory across specific state-action pairs to guide learning. Thus, while these methods are applicable, they do not address the unique challenges posed by bagged rewards, as they fail to account for the reward structure within the bags. Therefore, they are not ideal solutions for RLTF (Gangwani et al. 2020: Ren et al. 2021).
trade-off inherent in sparse rewards, while RLTF
In Fig. 6 we present additional results comparing our method with sparse reward baselines. Shaping Tessler et al. 2019 Hu et al. 2020) and LIRPG (Zheng et al., 2018) are described in the main paper. The Random Network Distillation (RND) [Burda et al., 2019b) method introduces an exploration bomus for deep reinforcement learning by measuring the prediction error of a randomly initialized neural network, which is a widely recognized and utilized approach for addressing sparse reward challenges. We include it as a baseline in our comparison of sparse reward methods for a more comprehensive evaluation.
These results demonstrate that, despite the superficial similarity between sparse rewards and bagged rewards, the internal structure of the rewards differs, leading to suboptimal performance of sparse reward baselines
21

---

Published in Transactions on Machine Learning Research (04/2025)
Table 4: Performance comparison across reward bag with various-length configurations over trials. In this table, "Short" refers to bags with lengths varying from 25 to 200, and "Long" denotes the setting with bag lengths from 100 to 500. The best and comparable methods based on the paired t-lest at the significance level 5% are highlighted in boldface.
Bag Setting
Environment
SAC
IRCR
RRD
HC
Shaping
LIRPG
RBT(ours)
3.21
269.12
2661.56
20.03
-596.02
-1407.15
5359.85
Ant-v2
(1.97)
(224.66)
(1675.65)
(46.35)
(178.75)
(504.59)
(129.28)
286.53
3275.05
2508.76
474.21
1358.92
287.38
3433.06
Hopper-v2
(72.51)
(44.68)
(690.07)
(66.93)
(640.58)
(114.65)
(96.89)
Short
5.92
10480.33
10382.80
3982.43
2010.72
1225.35
11073.88
H
(23.06)
(202.27)
(516.85)
(133.75)
(1085.36)
(162.88)
(181.43)
222.21
3840.30
3999.18
348.36
926.49
328.97
5198.09
W2
(66.72)
(666.24)
(561.49)
(174.28)
(100.37)
(107.69)
(225.50)
-115.10
215.38
2600.64
-0.53
-60.51
-2552.99
4897.50
Ant-v2
(138.93)
(92.94)
(1229.27)
(5.07)
(258.06)
(419.44)
(292.93)
360.36
3015.96
3089.52
652.28
729.28
325.84
3447.64
Hopper-v2
(118.06)
(408.08)
(333.23)
(92.41)
(338.56)
(50.25)
(83.02)
Long
-115.56
5944.17
8591.38
4563.99
-268.64
-571.35
10880.76
wa
(35.01)
(3421.46)
(1048.08)
(568,19)
(114.37)
(374.64)
(441.65)
251.11
3397.93
1221.37
497.27
779.43
281.27
4979.07
Walker2d-v2
(144.24)
(682.19)
(282.75)
(134.06)
(219.75)
(7.54)
(166.95)
in bagged reward environments. Conversely, as shown in the Section Compare with SOTA Methods of the main paper, our method performs well even in sparse reward environments, surpassing baselines specifically designed for sparse reward scenarios. This highlights the broad applicability and adaptability of our approach across different reward settings.
C
Additional Experimental Results
This section provides further analysis and insights through additional experiments to complement the main findings presented our main paper.
C.1 Experimental Result of Various-Length Reward Bags
In Table 4 we present the experiment results on various-leugth reward bags. The experiment depict in the table showcases the results of various methods applied across different environments with varying bag lengths of rewards, where bags
next to another as in the definition of RLBR. This experiment reveals that longer bags tend to degrade the performance of most methods. However, our RBT method appears to be less ivesitive to changes in bag length, maintaining robust performance even when the bag length is equal to the full trajectory. This result aligns with the result in the Section Compare with SOTA Methods in our main paper.
C.2 Preference on Handling Complex Delayed Reward Structures
The previous experiments were based on the assumption that bagged reward is the sum of immediate rewards, which aligns with prior works (Gangwani et al. 2020 Ren et al. 2021 Zhang et al., 2023). However, our approach does not make this assumption, indicating that the proposed method can effectively handle more complex reward structures. In this section, we experimentally validate this claim. Specifically, we included the following complex delayed reward structures computed over B
SumSquare:
The delayed reward is the sum of squared step rewards, placing more emphasis on larger rewards: R.o -iti abs(rr) Tt-
22

---

Published in Transactions on Machine Learning Research (04/2025)
The delayed reward is the square of the sum of step rewards, highlighting overall sequence performance: Reo abs Ei= (r) Ei=r- rt).
SquareSum:
Max:
The delayed reward is softmax-weighted sum of step rewards, giving more attention to the highest rewards: Roo = >t -ita-1 nyeeff
c"" Tt
where 3 is a scaling parameter that controls the sharpness of the softmax distribution.
The normalized scores shown in the figures computed by scaling the obtained rewards against the maximum possible rewards achievable in these bagged reward environments. For SumSquare (see Figure 7 our method consistently outperforms baseline methods across different environments, maintaining strong performance even as the bag length increases.
the SquareSum setting (see Figure [] the structure of bagged reward is more complex, as it emphasizes the cumulative effect of multiple steps rather than focusing on individual large rewards, making it more difficult attribute rewards specific actions. Despite this complexity, our method still holds clear advantage decreases with longer bag lengths. In the Max setting (see Figure g where rewards are sparse and only a few key steps contribute to the overall outcome, our method manages to learn effectively, showing that it can still identify and focus on the most critical instances in the sequence, even with limited reward signals. Under these conditions, methods relying on the sum-form delayed rewards assumption experience a significant drop in performance. This indicates that these met are heavily dependent on this assumption. is also worth noting that the HC method does not rely on this assumption, but its effectiveness limited to shorter delays. As the length of bag increases, its performance drops sharply.
most basclines,
performance slightly
Overall, our proposed method shows a strong ability to handle various bagged reward structures and adapt to increasing lengths, outperforming bascline methods in most environments. Morcover, these results highlight the importance of further exploring how to learn reward models and train policies in enviromments with complex bagged reward structure, making it worthy deeper investigation.
C.3 Architecture Sensitivity
Fig. 10 illustrates the sensitivity of our architecture to different imput sequence lengths during training (Seq len) and prediction lengths during reward relabeling (Relabel len) in the Ant-v2 environment. The learning curves, based on three independent runs with random initialization, show how varying these parameters affects the performance of the agent.
Although the configuration highlighted by the red box (Seq len of 500 and Relabel len of 100) demonstrates the best performance, the results also show that our proposed model is capable of learning effectively across various other configurations. This analysis underscores the importance of tuning imput sequence length during training and prediction length during reward relabeling for optimal performance, while also demonstrating the ability the RBT model to learn under different parameter settings, showcasing its flexibility and effectiveness reinforcement learning tasks.
23

---

Published in Transactions on Machine Learning Research (04/2025)
Ant
Norm Score
0.
0.4
23
50
100
200
500
Bag Lengthh
HalfCheetah
Norm Score
0.6
0.4
25
50
100
200
500
Bag Length Walker2d
Norm Score
0.6
0.4
0.2
23
50
100
200
500
Bag Length
fish-upright
0.8
Norm Score
0.6
200
Bag Length
swingu
Norm Score
0.4
0.2
23
50
100
200
500
Bag Length
SAC
URPG
HC
IRCR
RAD
RBT(ours)
Figure 7: Performance comparison on the SumSquare reward structure with six different bag length settings (5, 25, 50, 100, 200, and 500). The mean and standard deviation reported over six trials with different random seeds, measured across total of le6 lime steps.
24

---

Published in Transactions on Machine Learning Research (04/2025)
Ant
1.0
Norm Score
0.4
0.2
5
25
50
100
200
500
Bag Length
HalfCheetah
2.0
Norm Score
0.6
0.4
50
100
200
500
Bag Length Walker2d
Norm Score
0.0
0.4
50
100
200
500
Bag Length fish-upright
0.0
Norm Score
0.6
0.4
200
Bag Length
cartpole swingup
0.8
Norm Score
0.6
0.0
200
Dag Length
SAC
URPG
HC
IRCR
RRD
RBT/ours)
Figure 8: Performance comparison on the SquareSum reward structure with six diffcrent bag length settings (5, 25, 50, 100, 200, and 500). The mean and standard deviation are reported over six trials with different random seeds, measured across total of le6 time steps.
25

---

Published in Transactions on Machine Learning Research (04/2025)
Ant
Score
0.0
0.4
0.2
25
50
100
200
500
Bag Length
HalfCheetah
0.6
0.4
0.2
25
50
100
200
500
Bag Length Walker2d
Score
0.6
0.4
0.2
25
50
100
200
500
Bag Length
fish-upright
0.6
0.4
0.2
25
50
100
200
500
Bag Length
Score
0.6
0.4
0.2
25
50
100
200
500
Bag Length
SAC
URPG
HC
IRCR
AAD
RBT(ours)
Figure 9: Performance comparison on the Max reward structure with six different bag length settings (5, 25, 50, 100, 200, and 500). The mean and standard deviation are reported over six trials with different random seeds, measured across total of 1e6 time steps.
26

---

Published in Transactions on Machine Learning Research (04/2025)
Seq len 100, Relabel len 25
Seq len 200, Relabel len = 25
Seq len 500, Relabel len = 25
6000
6000
baq_len_100
9999
4000
4000
4000
Reward
Reward
Reward
8000
3000
5000
1000
1000
1000
200000
400000
600000
800000 "000000
290000
400000
600000
800000 1000000
200000
400000
630000
8:0000 1000000
Epoch
Epoch
Epoch
Seq 100, Relabel len
Seq 200, Relabel len
Seq
500, Relabel
6000
6000
6000
5000
5000
Reward
3000
3000
3000
2000
2000
2000
1000
600000
1000000
200000
600000
Epoch
Epoch
Epoch
5eq len 100, Relabel
100
len 200, Relabel len 100
Seg len
500, Relabel len 100
6000
6000
6000
5030
5000
5000
4000
4000
Reward
Reward
Reward
3000
3000
3000
2000
2000
2000
1000
1000
200000
400000
eсс0о0 8оо0сo 1000000
200000 400000 600000 800000 1000000
200000
400000 630000 sc0000 1000000 Epach
Epoch
Epoch
Figure 10: Learning curves on Ant-v2 with different length of input sequence in training (Seq len) and predict length during relabeling process (Relabel len), based on 3 independent runs with random initialization. Within each of the smaller graphs, the curves represent results from experiments with different bag lengths. Specifically, there are three bag lengths evaluated: 25, 100, and what labeled as 9999, which we interpret as proxy for trajectory feedback. The graph highlighted by the red box indicates our chosen parameter setting for the experiment, which is input sequence length of 100 and predict length of 500.
27

---

Published in Transactions on Machine Learning Research (04/2025)
C.4 Final Training Losses
Table 5: Final training losses at the trajectory level on the Ant-v2 environment with different bag lengths, showing the mean and standard deviation over three random seeds.
Bag Length (n)
50
100
200
500
9999
State Prediction Loss
0.003 + 0.002
0.023 + 0.001
0.031 + 0.007
0.032 = 0.013
0.056 + 0.036
0.053 = 0.027
0.032 + 0.011
Reward Prediction Loss
0.078 + 0.003
0.174 + 0.008
0.310 +: 0.049
0.769 = 0.061
1.342 + 0.146
3.066 + 0.529
6.195 + 0.406
Composite Loss
0,081 = 0.001
0.196 +: 0.008
0.341 + 0.056
0,801 = 0.068
1.399 +: 0.116
3.120 +: 0.520
6.227 + 0.396
The training losses vary with different bag lengths, reflecting how the model's learning process is influenced by the difficulty of reward redistribution and state prediction. To ensure a fair comparison, the estimated reward prediction loss, estimated state prediction loss, and estimated composite loss presented in Table 5 are all computed based on the entire trajectory. This mcans that for each trajectory, we first sum the reward prediction loss over all reward bags and the stale transition loss over all state-action transitions. The overall composite loss is then obtained by summing these two aggregated losses. The results show that the state prediction loss remains relatively stable across different bag lengths, suggesting that the model effectively learns environmental dynamics with minimal impact from the bag size. In contrast, the reward prediction loss increases significantly with longer bag lengths. This growth not only due to the increased difficulty of redistributing rewards over a larger set of instances but also stems from the inherent squared error accumulation in Equation 5 Even if the per-instance prediction error remains constant, the total loss scales with bag length because the squared error accumulates over all instances within the bag. For example, assurming constant instance-level error e and trajectory length 100, a short bag of length n = 5 results in an accumulated squared error of (100/5) x (5 x e)) = 500 x e², whereas a long bag of length n = 50 leads to a much larger accumulated error of (100/50) x (50 x e) = 5000 x x?. This explains why longer bags naturally lead to higher reward prediction loss, even when individual instance errors remain the same. Consequently, since composite loss is combination of state prediction loss and reward prediction loss, it is primarily driven by the increasing reward prediction loss al larger bag lengths. This suggests that the main challenge in learning under bagged rewards lies in accurate reward redistribution when fewer direct supervisory signals are available. Future improvements could focus on mitigating the impact of large bag lengths on reward prediction, such as incorporating regularization lechniques, hierarchical reward redistribution, or adaptive weighting in composite loss to improve overall learning efficiency.
Table 6: Normalized reward prediction losses in the Ant-v2 environment for different bag lengths, computed as the mean squared error between the sum of predicted instance-level rewards and the bagged reward, normalized by the square of the bag length. The table reports the mean over three random seeds.
Bag Length (n)
25
50
100
200
500
9999
Normalized Loss
1.56 x 10-5
6.96 x 10-6
6.2
10-6
7.69
10-6
6.71
10-6
6.13 10-6 6.20
10-6
We also present normalized reward prediction losses in Table
6 The normalization is performed as:
Normalized Loss (B) n²
(8)
where (B) represents the reward prediction loss for single bag, computed the squared difference between the sum of predicted instance-level rewards and the bagged reward. This normalization is motivated by the fact that (B) is mean squared loss, mcaning that its value naturally scales with n² due to the summattin inside the squared term. By dividing by n²² we obtain a loss that is comparable across different bag lengths. This adjustment allows for a fair per-instance comparison across different bag lengths. Interestingly, shorter bag lengths tend to have a higher normalized loss. Since the loss function aggregates squared differences within each bag, a possible reason for this is that having more bags means errors are accumulated and squared more frequently, leading to overall increase in per-instance loss.
28

---

Published in Transactions on Machine Learning Research (04/2025)
Theorem establishes that if the sum of redistributed rewards within each bag perfectly matches the bagged reward, the optimal policies in the BRMDP and the redistributed reward MDP are identical. However, in practical settings, learning reward model introduces approximation errors, meaning the redistributed rewards may not sum exactly to the bagged rewards. small non-zero loss in reward redistribution leads slight deviations in cumulative rewards along trajectories. While such deviations exist, our results Fig. 3 suggest that the learned policies still achieve performance comparable to those optimized with true rewards, indicating that the impact of small redistribution errors is limited. However, if the discrepancies were larger, they could potentially distort policy learning and lead to different optimization outcomes. In such cases, additional techniques, such methods for enhancing policy robustness to noisy rewards (Wang et al. 2020,, may be required to mitigate these effects. While our theoretical amalysis assumes perfect reward redistribution, this assumption holds approximately in practical implementations where the learned reward model achieves low prediction error. Investigating the precise impact of redistribution errors on policy learning is an important direction for future work.
C.5 Runtime and Computational Overhead
Table 7: Training runtime (in hours) for different methods.
Method
SAC
TRCR
Shaping
LIRPG
RRD
HC
RBT (ours)
Runtime (hours)
1
8
21
30
22
The Table 7 presents the training runtime for different methods on an NVIDIA GePorce RTX 2080 Ti GPU cluster with 8GB of memory. Only RBT (ours) and RRD utilize reward model, which introduces additional computational overhead, resulting in longer training times compared to methods that do not rely on a reward model. HC modifies the value function to adapt to sequence-level rewards, which also increases training time.
Notably, the primary factor contributing to the runtime of our proposed RBT method is not the training of the model itself but rather the relabeling of rewards for samples in the replay buffer. This process adds significant computational overhead, making it an area for potential optimization. A potential improvement is to use prioritized experience replay, where more informative samples are prioritized for relabeling instead of treating all samples equally, thereby improving efficiency while maintaining learning quality. Another approach to implement a cached reward update mechanism, where only a subset of samples is updated per iteration rather than recomputing rewards for all samples, significantly roducing computational costs. By incorporating these optimizations, the compulational efficiency of RBT has the potential to be significantly improved without compromising performance. By incorporating these optimizations, the computational efficiency of RBT has the potential to be significantly improved without compromising performance.
29


Reviews:
[{"review_id": "XLXiEH1Bx0", "reviewer": "Reviewer_BSYC", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Overall this is a nicely written paper which tackles a clearly defined problem simply but well.\n   1. The problem setting is well presented, the choices and limitations are clearly specified and experiments are of limited scope (i.e. Control Suite only) but well executed.\n   2. The proposed Transformer, which does both state prediction alongside reward prediction, is rather standard but the combination of Causal + Bidirectional layers does seem to work well and beats baselines.\n   3. I appreciated Figure 4, which demonstrated how well the transformer performed at predicting rewards, it was informative and clear.\n\nMy main issue with the paper is probably about generalizability and knowing if this would handle more complex domains.\n   1. The bag partition function G is assumed to be part of the environment. Unfortunately that seems like a strong requirement, and I can imagine that lifting this would make the problem setting much more complex (i.e. one would have to infer bag transitions, etc)\n   2. Most of the main text revolves around assuming that the bagged reward is the sum of intermediate rewards. This isn’t strictly necessary for their work though, as covered in Appendix D.2, but this makes this problem setting look a bit more simple than it really could be. It might be worth bringing this up earlier in the main text (e.g. when defining the bagged reward as the sum).\n   3. There is no strict presentation of generalisation ability across games / when domain transfer of the Transformer is expected, as far as I could see? All results present policies+reward functions trained and evaluated on environments independently, which is fair but again limited in scope.\n   4. How does the reward function transfer across environments? Do you expect it to?\n   5. The Control Suite is a simple and very smooth environment. Rewards are cyclical and simple to predict, the state space is also simple. How does the method handle more complex environments with partial observability and very sharp reward transitions? What about when exploration is difficult?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper tackles the problem of learning policies which can only observe delayed bagged rewards instead of instantaneous timestep rewards. They learn to invert the bagging transform (i.e. go back to instantaneous rewards) using a Transformer. They then apply off the shelf RL algorithms to obtain a policy. They show that their method is able to tackle simple control problems in the Control Suite.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Critical: (none)\n\nStrengthen:\n1. Discuss the limitations above more directly", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "u3n5ncjmqa", "reviewer": "Reviewer_voSh", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The experiments compare many baselines and are well-executed\n- The paper addresses a central issue of RL, and the proposed solution seems to be superior for dense and spare reward environments, which is an exciting result.\n\nWeaknesses:\n- The theorem and respective proof are not clear.\n- The paper lacks a limitation section that discusses caveats of the results and method.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes bagged-reward MDPs, which can be viewed as an advanced-reward shaping approach. In particular, the authors propose a transformer to distribute bagged rewards among the chosen action state pairs. The approach is evaluated against other more naive ways to distribute the reward and for different bag time horizons on MuJoCo tasks. This shows the clear superiority of using the transformer to reshape over the other approaches. Additionally, the authors perform an ablation study for their transformer architecture demonstrating that all proposed components (bi-directional attention and state estimation) are necessary for the best performance.\n\nThe paper is well written and the figures and tables support the understanding well.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- Theorem 1:\n  - My understanding of the theorem is that if the sum of the redistributed reward of one bag is equivalent to the bag reward, the MDP and BRMDP are equivalent. If that is correct, this is not clear from the theorem, please rephrase or explain\n  - I could not follow the proof logic - if my assumption is correct the last equation in the proof would suffice in my opinion. I would include the shorter proof in the main text or give an abbreviated proof there.\n  - Also the wording is not identical for Theorem 1 in the text and in Appendix A, this should be corrected\n\n- Limitation section - I see several undiscussed points that could be answered by a limitation section, additional experiments or revised related work depending on your preference:\n  - What training is necessary for the transformer (how long, how many trajectories are used, what are the final losses $\\mathcal{L}_{bag}$, $\\mathcal{L}\\_{r}$ and $\\mathcal{L}\\_{s}$ ..)? Do you train a new one for each task? \n  - How do you ensure that the sum of the redistributed reward is equal to the bagged reward?\n  - Why did you not compare to vanilla SAC in your experiments? This would be the minimum that should be achievable and a good reference for all other approaches.\n  - While your approach addresses many issues of reward engineering, it still assumes a reward function (sparse or dense) is available. This can still be difficult to obtain in the first place. What do you think about extending your approach with rewards that are automatically generated from task specifications [1-3]? \n\n- The contribution paragraph is not too clear. Pleas separate contributions and your solution approach more clearly. \n- The origin of the bagged rewards was not too clearly introduced. I would appreciate a revised experimental setup that more clearly states what is given and what is calculated.\n\nReferences:\n\n[1]  Rodrigo Toro Icarte, Toryn Q. Klassen, Richard Valenzano, and Sheila A. McIlraith. 2022. Reward Machines: Exploiting Reward Function Structure in Reinforcement Learning. J. Artif. Int. Res. 73 (May 2022). https://doi.org/10.1613/jair.1.12440\n\n[2]  LTL and Beyond: Formal Languages for Reward Function Specification in Reinforcement Learning. Alberto Camacho, Rodrigo Toro Icarte, Toryn Q. Klassen, Richard Valenzano, Sheila A. McIlraith. Proceedings of the Twenty-Eighth International Joint Conference on Artificial Intelligence Understanding Intelligence and Human-level AI in the New Machine Learning era. Pages 6065-6073. https://doi.org/10.24963/ijcai.2019/840\n\n[3] Le, X. B., Wagner, D., Witzman, L., Rabinovich, A., & Ong, L. Reinforcement Learning with LTL and $\\omega $-Regular Objectives via Optimality-Preserving Translation to Average Rewards. In The Thirty-eighth Annual Conference on Neural Information Processing Systems.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "WUL9Kcnwwz", "reviewer": "Reviewer_u6fd", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The submission is well-written and easy to follow. \n\nThe results are thorough with both theoretical and thorough empirical analysis. The improvement of the proposed method is well demonstrated. \n\nAccording to my knowledge, the proposed method is novel and tackles an important problem, where the reward is defined on a trajectory in a non-Markov and not-necessarily additive manner.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The submission considers decision-making where the reward is sparse and defined on trajectories instead of one single state-action pair. To tackle this problem, the authors first propose to formulate the problem as an MDP with bagged reward and connect this formulation withe an ordinary MDP but with properly redistributed reward. In other words, the key is to properly redistribute the reward. Therefore, the authors propose a novel reward redistribution method built transformers. \n\nBoth theoretical and empirical analysis are provided, demonstrating the advantages of the proposed method.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I do not have any suggested adjustments.", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper addresses an important problem in reinforcement learning—credit assignment under delayed or sparse rewards—by introducing a novel reward redistribution method using bidirectional attention mechanisms. The reviewers unanimously recommended acceptance, citing the paper's strong theoretical foundation, clear presentation, and empirical rigor. However, they also raised some valid concerns about generalizability, scope, and assumptions. Below is a summary of these points and how they were addressed:

Generalizability Across Domains:
Reviewers noted that experiments were limited to simple environments (e.g., Control Suite) without demonstrating cross-domain generalization. The authors acknowledged this limitation and discussed it in the revised manuscript's "Limitations" section. They proposed future work on pretraining across environments or incorporating meta-learning techniques to enhance transferability.

Assumption of Known Bag Partition Function:
While the assumption simplifies the problem, reviewers suggested that lifting this constraint would make the approach more broadly applicable. The authors justified their choice by aligning it with practical scenarios but acknowledged this as a limitation and discussed potential extensions involving latent variable modeling.

Complex Environments:
Concerns were raised about how the method would perform in more complex settings with partial observability or sharp reward transitions. The authors responded by highlighting their method's adaptability through sequence-based modeling while acknowledging challenges in highly stochastic environments.

Scope of Reward Aggregation Assumption:
Reviewers pointed out that assuming bagged rewards are sums of intermediate rewards oversimplifies the problem setting. The authors clarified that this assumption is specific to their method but not fundamental to RLBR as a whole, revising Section 3 for better transparency.

These revisions adequately address the reviewers' concerns without requiring substantial additional work from the authors before publication.

The Action editor encourages the authors to connect with use cases in embodied and digital agents as complex tasks naturally admit to decomposition, hence bagged rewards.
Decision: Accept (poster)