A Snapshot of Influence: A Local Data Attribution Framework for Online Reinforcement Learning

Thanks for the detailed, thoughtful and constructive feedback! We are grateful that the reviewer appreciates the contributions of our work, and address their comments in what follows.

PPO+Adam

Thanks for raising this important point! We would like to clarify that while our IIF framework is developed for SGD, it can be used to accelerate training when the optimizer is Adam. In particular, we already used Adam in BipedalWalker (Fig. 6, last environment) and RLHF (Fig. 7) for both standard training and IIF (details in Appendices B.2 and B.3). In these cases, following [1], we train the model with Adam but compute the influence scores as if the model is trained with SGD, i.e., using the same formula below line 156. This approach is not perfect (we have explicitly listed it as a limitation & future work in Sec. 6), but it has already achieved substantial gain!

For the first 5 environments in Fig. 6, we used SGD for both standard and IIF (details in Appendix B.2), since it already performs well in these environments. Following the reviewer’s comment, we conducted additional experiments on two of these environments using Adam. As the rebuttal format is text-only, we present the sample efficiency and runtime results in the table below, following Fig. 6(b).

Table: Improvement of IIF compared to standard training in sample efficiency and runtime. Optimizer = Adam; learning rate = 3e-4, 1e-3 for MiniGrid and LunarLander, respectively.

The results show that IIF delivers a clear and substantial benefit regardless of the choice of optimizers or environments, confirming our framework’s generality. One observation is that IIF gains less with Adam compared to SGD in MiniGrid, whereas the trend is reversed for LunarLander. This is partly because Adam significantly speeds up training compared to SGD in MiniGrid (and thus reduces the room of improvement), but less so in LunarLander. We will include these new results as well as discussions on the optimizers in the revision.

RL terminologies

1. Online RL vs. On-policy RL: Thanks for highlighting the important distinction. In the revision, we will 1) add a paragraph in Sec. 3 to explicitly state that the core challenge we address—the circular dependency between data and model—is a fundamental characteristic of the entire online RL paradigm, while the solution we present and evaluation we perform is specific to PPO, an on-policy algorithm; 2) discuss how our framework can generalize to other on-policy and off-policy RL algorithms (see “Generalizability / extensions”)

2. Record vs. transition: We chose the term “record” because it more accurately reflects the unit of our data attribution method. In our framework, a record is defined as $z_i = (s_i,a_i,r_i,\log \pi_i, v_i,\hat A_i)$, which encompasses not only a raw transition but also computed values like the log probability and the advantage estimate, which are essential for influence calculation. Moreover, “record” more accurately captures the atomic unit in RLHF, which is a (prompt, response) pair. That said, we appreciate the reviewer’s suggestion and will clarify how a “record” is related to a “transition” in the revision.

Generalizability / extensions

By design, our framework can generalize to other on-policy and off-policy algorithms with straightforward adaptations. The core principle of measuring influence via gradient similarity is universal. For on-policy methods, this primarily involves replacing the PPO gradient (the $\mathcal{L}^{\text{PPO}}$ term under line 156) with the loss gradient of the alternative algorithm. For value-based off-policy algorithms like DQN, we additionally need to change the target function to be the Bellman error. We will add these discussions in the revision.

Clarification questions

1. Hyperparameter $p$: We agree with the reviewer that an additional hyperparameter entails additional tuning; however, from our ablations in Appendix C.6, we have demonstrated that any level of filtering can improve performance over standard training, and that a value between 12.5% and 50% works well across settings.

2. TracIn calculation: We use $f^{\text{return}}$ to explain the two gradient terms in TracIn. The first is the target gradient which does not depend on $z_i$’s ($\nabla_\theta f^{\text{return}}(\theta_j)$), computed on a validation set and represents the desired direction for policy improvement. The second is the per-record gradient ($\nabla_\theta \mathcal{L}^{\text{PPO}}(\theta_j, z_i)$), computed for a single training record $z_i$ and captures its learning signal. A positive dot product between these two vectors indicates the record is helpful for policy improvement, whereas a negative score indicates it is harmful.

3. Role of target functions: The two target functions have different use cases. $f^{\text{action}}$ is mainly for diagnosis—understanding why the agent takes a specific action (e.g., abnormal / harmful) at a specific state. On the other hand, $f^{\text{return}}$ assesses contribution to overall performance. This makes it suitable for both analysis (as in Sec. 4.1) and algorithmic policy improvement. We believe it is not suitable to use $f^{\text{action}}$ to guide RL training because it is not tied to a universally desirable outcome. Using it would require an external "oracle" to determine which targeted behaviors to reinforce at each specific state, which, however, is the very problem that RL is designed to solve.

IIF in sparse-reward settings

Our new results on MiniGrid show IIF's benefit is smaller with Adam, suggesting part of the gain may stem from compensating for SGD's inefficiencies. Nevertheless, we hypothesize that IIF could potentially be more useful in sparse-reward settings. The noisy and misleading advantage estimates common in such environments are precisely what our method identifies as "harmful records", making them ideal for filtering via IIF. Investigating the joint effects of environment and optimizer on IIF’s performance is an important direction for future work.

Final note

If you find our response satisfactory, we kindly ask you to increase your rating to support our work. Thanks again for your time and effort in improving our work!

References

[1] Wang, Jiachen T., et al. "Data shapley in one training run." ICLR 2025.

Thanks for the detailed, thoughtful and constructive feedback! We are grateful that the reviewer appreciates the contributions of our work, and address their comments in what follows.

Statistical significance of experimental results

We appreciate this critical feedback, and have performed additional experiments in three environments to address the concern. Specifically, we follow [1] to use 5 random seeds and report the mean $\pm$ standard error. As the rebuttal format is text-only, we present the sample efficiency and runtime results in the table below, following Fig. 6(b).

Table: Improvement of IIF compared to standard training in sample efficiency and runtime. Experimental configs are the same as in Fig. 6(b).

The results show that IIF provides substantial improvements in sample efficiency and runtime, and these gains are statistically robust. We will update all figures and tables in the revision using 5 random seeds.

Q1: Role of target functions

The two target functions have different use cases. $f^{\text{action}}$ is mainly for diagnosis—understanding why the agent takes a specific action (e.g., abnormal / harmful) at a specific state. On the other hand, $f^{\text{return}}$ assesses contribution to overall performance. This makes it suitable for both analysis (as in Sec. 4.1) and algorithmic policy improvement. We believe it is not suitable to use $f^{\text{action}}$ to guide RL training because it is not tied to a universally desirable outcome. Using it would require an external "oracle" to determine which targeted behaviors to reinforce at each specific state, which, however, is the very problem that RL is designed to solve.

Q2: Target function $f^{\text{return}}$

Thanks for this excellent question! Our goal is to measure a sample's influence on the agent's true objective (maximizing return), so we use $f^{\text{return}}$ as a clean and direct target. The PPO loss, in contrast, is a complex surrogate designed for stable training—it combines policy, value, and entropy terms. Attributing influence to this composite loss would be difficult to interpret; it is unclear whether a sample is deemed influential because it improves the policy, refines the value estimate, or simply affects exploration via the entropy term. By decoupling the optimization algorithm from the evaluation objective, our approach allows us to ask a much more direct and meaningful question: "Does this sample help the policy achieve higher expected returns?"

Clarity

1. Figs. 1, 3, 4: We appreciate the comment and will revise the figures and their captions for better readability.

2. Online RL vs. On-policy RL: Thanks for highlighting the important distinction. In the revision, we will 1) add a paragraph in Sec. 3 to explicitly state that the core challenge we address—the circular dependency between data and model—is a fundamental characteristic of the entire online RL paradigm, while the solution we present and evaluation we perform is specific to PPO, an on-policy algorithm; 2) discuss how our framework can generalize to other on-policy and off-policy RL algorithms (see “Generalizability / extensions”)

3. Statements (L111, L116):

• L111: Retraining-based methods require training the model once for each of the $N$ records being evaluated. This is computationally expensive in any setting and particularly prohibitive in RL, where a single training run involves a long sequential process that can take hours, making $N$ separate runs infeasible.

• L116: As illustrated in Figure 2, a training record has two influence pathways: direct parameter updates (green arrows) and subsequent data generation (red arrows). Any attribution method that ignores the second channel captures only a partial effect, yielding scores that likely deviate significantly from the true effects.

4. The action target function: This target function does not require an expectation because it is designed as a diagnostic tool for a specific behavior of interest (taking an action $a$ at a state $s$). In contrast, $f^{\text{return}}$ concerns maximizing the expected return over a distribution of states, thus requiring a validation dataset.

5. Role of target functions: See “Q1: Role of target functions”

6. Generalizability / extensions: By design, our framework can generalize to other on-policy and off-policy algorithms with straightforward adaptations. The core principle of measuring influence via gradient similarity is universal. For on-policy methods, this primarily involves replacing the PPO gradient (the $\mathcal{L}^{\text{PPO}}$ term under line 156) with the loss gradient of the alternative algorithm. For value-based off-policy algorithms like DQN, we additionally need to change the target function to be the Bellman error. We will add these discussions in the revision.

7. Clarifications of L137, L142:

• L137: Using reference distributions that vary across runs is a direct consequence of our contextual, local framework. The goal is not to identify universally "good" data, but rather to answer a more immediate question: "For the agent at its current stage of training, which of its recent experiences were most helpful or harmful for the next training update?"

• L142: We sample trajectories using the reference policy $\pi^{\text{ref}}$, i.e., $a \sim \pi^{ref}(s)$.

8. Benefit of dynamic reference: This is related to the previous question (L137). Using a dynamic reference is a deliberate choice to measure progress that is contextual to the agent's current state. A fixed, off-distribution reference can be misleading due to distribution mismatch; for instance, a novice agent gains little from a dataset sampled from an expert policy which covers states it never reaches.

   To demonstrate this empirically, we performed additional experiments of single-round interventions (as in Sec 4.3, lines 221-223) using a fixed reference dataset from a well-trained policy. The results show that our intervention consistently improves the target function ($\Delta f^{\text{return}} > 0$) but degrades actual test performance ($\Delta \text{return} < 0$) except for the few final rounds. This validates our claim that an evolving, on-policy reference is crucial to ensure effective intervention.

Experimental design and analyses

1. Random seeds: See “Statistical significance of experimental results”

2. Measure of variance: Thanks for the thoughtful comment; we will update the figure captions to improve clarity. The boxplots in Fig. 5 show the full data distribution (median, IQR, and extent) over 3 random seeds. The shaded areas in Fig. 6 represent the standard error, used to visualize confidence and assess the statistical significance of the differences between the methods [1].

Related work

Thanks for pointing out the relevant literature! We would like to note that we have already compared our method with an adaptation of Prioritized Experience Replay [2] in the on-policy setup (see Appendix C.5). In the revision, we will include a more thorough discussion of this line of work, and also discuss [3] in Sec. 4.1.

Final note

If you find our response satisfactory, we kindly ask you to increase your rating to support our work. We are happy to engage in any follow-up questions. Thanks again for your time and effort in improving our work!

References

[1] Henderson, Peter, et al. "Deep reinforcement learning that matters." Proceedings of the AAAI conference on artificial intelligence. Vol. 32. No. 1. 2018.

[2] Schaul, Tom, et al. "Prioritized experience replay." arXiv preprint arXiv:1511.05952 (2015).

[3] Ilyas, Andrew, et al. "Are deep policy gradient algorithms truly policy gradient algorithms." arXiv preprint arXiv:1811.02553 (2018).

Thank you for the response.

Random seeds

We point out that our experiments follow exactly the same practice as in [1]. Specifically, in their "Experimental Analysis" section, the authors state: “To ensure fairness we run five experiment trials for each evaluation, each with a different preset random seed (all experiments use the same set of random seeds)...All results (including graphs) show mean and standard error across random seeds.” While they also discuss alternative approaches such as power analysis, their own empirical evaluation relies on five seeds.

Furthermore, we wish to clarify a potential misinterpretation. Based on our reading, [1]’s primary caution is not against using five seeds per se, but against the lack of reproducibility when comparing results derived from different sets of seeds. Our work strictly adheres to their recommendation by using a consistent set of seeds for all comparisons.

Ultimately, we believe the statistical significance of our method's improvement is clearly demonstrated by the results. We invite all reviewers and the AC to examine Figure 6(b) as well as the new results in our response to Reviewer G83y, where the mean performance and standard error clearly demonstrate a significant gain over the baseline.

Complexity of the environment

Reviewer 6kkY has already acknowledged that our response and newly added experiments resolved their question.

Application to PPO

As stated in the abstract and throughout the paper, our work is intentionally focused on the “widely used Proximal Policy Optimization (PPO) algorithm”. Our goal is to provide a deep and rigorous analysis for a highly relevant, state-of-the-art algorithm, rather than a shallow evaluation across multiple methods.

While our framework can be easily extended to other online RL algorithms (as discussed in our previous response), we believe that demonstrating this on older methods like TRPO or DQN would add limited value. PPO’s prevalence makes it the most impactful testbed for a foundational study. We note that this focused approach is common in the literature. For example, the foundational work on influence function [2] introduced its methodology for understanding black-box predictions but focused its evaluation solely on vision tasks. This seminal work has since been successfully applied to numerous other domains. Similarly, we aim to establish a robust method on a critical algorithm, paving the way for broader application.

Clarity

The reviewer mentioned “abundance of clarity concerns” and unsure “whether they will be fixed”. Could the reviewer provide concrete examples of what remains unclear so we can direct our efforts effectively? We respectfully note that other reviewers found the paper’s writing to be a strength (e.g., “well written and easy to follow” and “good presentation”).

References

[1] Henderson, Peter, et al. "Deep reinforcement learning that matters." Proceedings of the AAAI conference on artificial intelligence. Vol. 32. No. 1. 2018.

[2] Koh, Pang Wei, and Percy Liang. "Understanding black-box predictions via influence functions." International conference on machine learning. PMLR, 2017.

Thanks for giving us a chance to better understand your thoughts and concerns. Below, we address them and propose changes which we will make in the revision.

• Final performance CI. We concur with the reviewer on rigorous statistical analysis. To this end, we computed the 95% confidence interval (CI) for the performance gain of IIF over the standard baseline across 5 random seeds (half-width = $t_{0.975,4} \times SE$ = $2.776\times SE$). The results confirm a statistically significant improvement: the 95% CI is [22.54, 130.52] for LunarLander and [24.40, 75.99] for BipedalWalker.

  Proposed change: 1) L277, We will change “IIF’s final performance exceeds standard training in almost every environment” to “IIF’s final performance exceeds standard training in challenging tasks”. 2) We will include 95% CI for all reported results in Figs. 5, 6, 7 in the revision.

• Online RL vs on-policy RL. We argue that extending our framework to off-policy RL is not merely about deriving similar attribution terms; the scores are meaningful. As noted earlier, the core of our framework is to ask: At the agent’s current stage of training, which data were most helpful or harmful for the next update? This principle applies regardless of whether data is on-policy or off-policy, recent or old; attribution aims to determine usefulness for learning. One important caveat here is that in on-policy RL, current training data can serve for validation, while in off-policy RL, to measure progress relevant to the current agent, we must sample fresh data from the current policy to form the validation dataset, thus sacrificing the sample efficiency of our framework.

  Proposed change: We will add a subsection (Sec. 3.2) to discuss how our framework can generalize to other on-policy and off-policy RL algorithms, with potential caveats (which could also serve as interesting future work) as described above.

We also detail the changes aimed at improving clarity, thanks to the reviewer’s suggestion:

• L111: We will change “Retraining-based methods are too costly and thus inapplicable” to “Retraining-based methods require training the model once for each of the records being evaluated, which is computationally expensive in any setting and particularly prohibitive in RL.”

• L116: We will change “Technically, one can compute influence scores ignoring this extra channel of influence, but … errors that may compound exponentially over rounds” to “If we compute influence scores using the original formulas from standard supervised learning, they capture only the impact on parameter updates, ignoring the extra channel of influences through future data generation. As a result, the scores may deviate significantly from the true influence we seek to measure.”

• L142: We will modify the equation to ​$f^{\text{return}}(\theta)=\mathbb{E}_{\tau\sim \pi^{\text{ref}},(s,a)\sim\tau}[\log\pi(a\mid s)\hat A^{\text{ref}}(s,a)]$.

• L149: We will change “This provides a stable and relevant reference that evolves with training progress, which is crucial for meaningful evaluation and thus attribution as the policy improves.” to “This is a key design choice of our contextual framework, which enables us to ask: For the agent at its current stage of training, which recent experiences were most helpful or harmful for the next update? Unlike a fixed, off-distribution reference that may provide misleading signals due to mismatch with the agent’s current state, our dynamic reference evolves with training, providing a stable and relevant basis for meaningful evaluation and attribution.”

• L154: We will add a remark below L154 to explain the use cases of these two target functions: “The two target functions have different use cases. $f^{\text{action}}$ is mainly for diagnosis: understanding why the agent takes a specific action at a specific state. On the other hand, $f^{\text{return}}$ (Sec. 4.2) assesses contribution to overall performance, which makes it suitable for both analysis (Sec. 4.1) and algorithmic policy improvement (Sec. 5).”

• Fig. 1: We will lead the caption with “An Illustration of on-policy RL’s alternating learning cycle (Sec. 2.1) and our local data attribution framework (Sec. 3.1)”, to clarify what the figure aims to convey and where the detailed descriptions are located. We will simplify the content in the figure and highlight the most important components for ease of understanding.

• Figs. 3 & 4: We understand that a lot of information required to understand these two figures are present only in the text following the plot. We will change the order of presentation, and also expand the content in the caption.

We are grateful that the reviewer upholds high standards for our work and are happy to engage in further discussions if they’d like to suggest other changes.

Thanks for the positive review and insightful questions! We are grateful that the reviewer appreciates the contributions of our work, and address their comments in what follows.

Results and analyses for other games

Following the reviewer’s suggestion, we conducted new case studies on two complex environments in the early-to-mid stage of training.

1. In Pong (Atari), we found that bottom records filtered by IIF consist of uninformative transitions (the ball being out of play or already moving away from the agent) that receive (inaccurately) high advantage estimates. By filtering out these samples, IIF achieves significant improvement in training efficiency.
2. In BipedalWalker (Robotics), our analysis revealed bottom records where the agent was incorrectly penalized with a large negative advantage for executing a successful recovery move (e.g., applying corrective torque with a deeply bent knee (~35°) during landing or push‑off).

These results show that 1) bottom records feature inaccurate advantage estimates; 2) IIF is effective, holding generally across different environments. We will include these new analyses and results in the revision.

Extension to other attribution algorithms

Yes, our framework is general and compatible with any gradient-based data attribution algorithm. The three core components of our framework (entity of attribution, target function, attribution method) are modular and can be designed independently. This allows for the straightforward integration of other attribution methods like influence functions [1] or SGD-influence [2].

We selected TracIn for this study primarily for its conceptual simplicity and empirical efficiency. We believe investigating the interplay between our framework and other attribution algorithms is a promising direction for future work.

Gradients of selected layers

Thanks for the thoughtful suggestion. In our work, influence is computed using the gradient with respect to the full set of model parameters.

While using gradients from selected layers can reduce computational costs, as explored in prior work [3,4], this approximation can introduce bias to the attribution scores (e.g., see [5]). Our work instead leverages the ghost dot product technique from [5]. It enables highly efficient influence calculation using the full gradient (and thus avoiding the bias of layer selection), while adding negligible computational overhead to standard training. A detailed runtime analysis is in Appendix C.7 (Table 3), where we show that the overhead from our full-gradient influence calculation is minimal.

Final note

We are happy to engage in any follow-up questions. Thanks again for your time and effort in improving our work!

References

[1] Koh, Pang Wei, and Percy Liang. "Understanding black-box predictions via influence functions." International conference on machine learning. PMLR, 2017.

[2] Hara, Satoshi, Atsushi Nitanda, and Takanori Maehara. "Data cleansing for models trained with SGD." Advances in Neural Information Processing Systems 32 (2019).

[3] Pruthi, Garima, et al. "Estimating training data influence by tracing gradient descent." Advances in Neural Information Processing Systems 33 (2020): 19920-19930.

[4] Yeh, Chih-Kuan, et al. "First is better than last for language data influence." Advances in Neural Information Processing Systems 35 (2022): 32285-32298.

[5] Wang, Jiachen T., et al. "Data shapley in one training run." ICLR 2025.

Thanks for the positive review and insightful questions! We are grateful that the reviewer appreciates the contributions of our work, and address their comments in what follows.

Explanation of the performance gain

Filtering based on data attribution scores allows us to remove samples with harmful learning signals. As summarized by Reviewer 6kkY: “The gradient direction of the value/policy function is the direction for improvements. If a sample loss function gradient doesn't align well with the above direction, this sample will hinder optimization… Thus, it should be removed.” This principle is powerful because attribution scores are directly tied to the target function we aim to optimize (e.g., expected return), leading to a highly targeted filtering process.

Beyond this general insight, our work develops a RL-specific interpretation of this effect. In particular, Sec 4.1 shows that the bottom records typically feature inaccurate advantage estimates. Removing them thus purifies the learning signal. This gives a more concrete, though partial (as noted in Appendix C.4) explanation to why filtering in RL is helpful.

Difficulty-based filtering

Thanks for this insightful point. To the best of our knowledge, difficulty-based filtering (e.g., pass@k) is primarily used to improve LLM Reasoning (RLVR), with GRPO being the most commonly adopted RL algorithm. In GRPO, advantages are normalized across multiple responses for the same prompt. Consequently, for an extremely easy or hard prompt where all responses receive similar rewards, the resulting advantage estimates are near-zero, providing little learning signal.

However, this logic does not apply to PPO. We demonstrate this with a new experiment in RLHF where we used reward as a proxy for difficulty and filtered records receiving top and bottom rewards. This heuristic performed worse than random because it systematically removed data with both highest and lowest influence scores, thereby harming the learning process. This finding aligns with our results in Appendix C.5 for traditional RL, where an analogous heuristic using TD error as a proxy for difficulty also proved ineffective. Therefore, our evidence shows that while valid for GRPO, difficulty-based filtering is an ineffective heuristic for PPO. We will add this discussion to the revision.

Extending our framework to RLVR with GRPO is a promising direction that we are actively working on. However, due to the many technical challenges involved, we believe this is best suitable for a dedicated follow-up work.

More applications of our framework

Thanks for this thoughtful question. In our paper, we have demonstrated that IIF achieves a significant improvement in training efficiency (4x reduction in wall-clock time) and final reward, which we believe is of sufficient interest to RLHF practitioners. Additionally, we believe our framework holds the potential in a few other applications:

1. Debugging and improving the reward model: A core challenge in RLHF is that the reward model is often imperfect. Our framework can identify training records that are “harmful” to the RLHF objective and trace this back to the flaws in the reward model—e.g., over-rewarding sycophantic or verbose but unhelpful responses. Practitioners can then use this set of records to debug and recalibrate the reward model, creating a cycle of iterative improvement.

2. Active prompt curation for efficient alignment: The efficiency of RLHF depends heavily on the quality of prompts used for generation. Our framework can be used to identify the set of prompts that provide the most value for policy improvement at any given stage of training. This would allow practitioners to strategically curate or up-sample the most informative prompts, leading to more targeted and efficient alignment in subsequent training.

Final note

We are happy to engage in any follow-up questions. Thanks again for your time and effort in improving our work!