Beyond Simple Sum of Delayed Rewards: Non-Markovian Reward Modeling for Reinforcement Learning

1. Discussion on the non-causality of weight calculation in Equation (7) is required.

• 1-1. The proposed architecture in Figure 2 uses a causal transformer for query, key, and $\hat{r}$. However, in equation (4), I do not understand why $w$ incorporates state-action pairs up to $i+n_i-1$, not $t$ as in the case of the query, key, and $\hat{r}$, which means that the learned instance-level rewards include future information. However, is it valid to assume that the current reward includes future information?
• 1-2. I recommend using causal mixing in equations (6) and (7) (i.e., use $k_{t'}$ for $t' \leq t$) and report the result as a part of the experiment section. (In this case, the weight calculation in (7) can adopt averaging over $i \geq t$ for each $t$, not summing over all $i$).
• 1-3. Plus, there is no $t$ in equation (7). I guess the authors unintentionally drop a subscript of $t$ by mistake.

2. Ablation study is required on core components: non-Markovian reward modeling and weighted sum.

• 2-1. The authors should report results by ablating the weighted sum part (e.g., use uniform weight instead) for both the sum-form and general composite form.
• 2-2. The authors should report results by ablating non-Markovian reward modeling for both sum-form and general composite form that can be done by (i) ablate causal modeling of query and key, (ii) ablate causal modeling of $\hat{r}$, and (iii) ablate both. Here "ablate" means not using the causal transformer and using MLP stepwise.
• 2-3. Is there any specific reason why equation (5) uses linear transformation, not a non-linear one? Can non-linear compression improve performance?

3. The novelty is not significant compared with a baseline (namely RBT)

• One of the proposed algorithm's two key components, modeling non-Markovian reward, appears to mirror the RBT algorithm closely. A main difference is the adaptive weight learning using softmax operation, but RBT also adopts softmax operation for the instance-level reward calculation. The authors should clarify the specific innovations over RBT more explicitly to guarantee their novelty.

4. Readability should be improved.

• 4-1. The authors should improve the motivation for the delayed reward setting. Is there any concrete example other than the human annotation perspective?
• 4-2. Section 2.3 should be included in Section 3 since CoDeMDP is the authors' proposed concept.
• 4-3. Placing the pseudocode in the main body can help readers' understanding for the algorithm.

5. I recommend increasing the number of seeds to more than 3 (at least 5).

Similarity with previous work:

This paper is extremely similar to Reinforcement Learning from Bagged Reward [1]. Both algorithms use transformer networks for the instance-level (per time step) reward prediction, where the attention mechanism operates on the sequences corresponding to the rewards. The algorithm's (and the pseudocode's) training process is also very similar.
In addition, both works assume non-Markovian rewards, even though the authors explicitly mention in the text that methods in this domain assume Markovian rewards.

The authors should clarify how their work differs from [1]. [1] https://arxiv.org/pdf/2402.03771

Missing related work:

The authors should mention the sub-field of preference-based RL, which can learn non-markovian reward models from sequences of (state, action) pairs, especially because the authors mention human feedback and evaluation in the introduction.

Unclear on Experimental Design Choices:

1. The authors decided to use 6 different reward delay lengths but elected not to run experiments in the truly sparse setting, where the agent receives a reward at the end of the episode. This seems surprising because the sparse reward setting is a commonly studied problem.

2. Why did you include these specific composite delayed reward structures, SumSquare, SquareSum, Max?

3. SumSquare is supposed to be the sum of the squared step rewards, but why is the formula the sum (abs(r)*r)? Shouldn’t it be (r2)? If the reward is negative, then abs(r)*r is not the same as (r2)? The same issue occurs in the SquareSum. Can the authors explain this?

Overclaiming on experimental results:

The authors claim their algorithm consistently outperforms baseline methods across different environments (L427). However, reviewing Figure 3, CoDeTr seems comparable to the other baselines. It would be helpful if the authors clearly specified in how many experiments their algorithm outperformed all baselines. In addition, running a statistical analysis (i.e., t-test) on these results would provide further support for any claims.

While the paper is well-structured and tackles an interesting challenge, I have some reservations about the importance of its contributions, mainly due to these two points:

1. The problem of non-Markovian or weighted delayed rewards feels somewhat artificial. Beyond the rather vague examples in the introduction, it’s unclear in what specific tasks or RL environments this problem would naturally occur. Without clearer context, the paper seems to be addressing a problem it has introduced.

2. CoDeTr decomposes delayed rewards into per-transition rewards, but these per-transition rewards are still non-Markovian, and this limitation seems unaddressed.

Due to the first point, the practical relevance of the proposed solutions is unclear, and the second point raises doubts about its theoretical value.

That said, I would consider recommending acceptance if these concerns are addressed in the authors' response.

[W1 & Q1] Reward and weight validity

Without additional constraints, it is unclear if $r$ and $w$ outputs genuinely reflect rewards and weights. Since $R_{co}$ is a weighted sum of $r$ and $w$ outputs, additional experiments could help clarify their interpretive accuracy.

[W2 & Q2] Performance in environments without actual per-step reward

The experiments are conducted on environments with known per-step rewards, artificially accumulated as composite rewards. Testing on complex environments with no per-step reward signals, such as robotics tasks with task-completion as the final reward or RLHF for language models with only a final reward, would enhance the assessment of CoDeTr’s performance and applicability in broader contexts.

[W3 & Q3] Impacts of model expressiveness

In Figure 4, the best two methods appear to be CoDeTr (proposed) and RBT. They both use a full-sequence Transformer model for credit assignment, while the other methods use a two-layer MLP. The performance of CoDeTr and RBT is also very close. The full-sequence Transformer model takes in the entire sequence as input and is more powerful. However, the MLP only takes in the current Markovian state/action as input and is less expressive.