Which Rewards Matter? Reward Selection for Reinforcement Learning under Limited Feedback

TL;DR

The paper addresses how to choose which states to label with rewards when feedback is limited, formalizing reward selection for reinforcement learning from limited feedback (RLLF) in offline datasets. It introduces a training-frame that evaluates strategies with an evaluator $\\Xi$ and analyzes both training-free heuristics and training-based optimization (brute-force, sequential-greedy, and ES) to identify the most informative reward labels. Empirically, it shows that effective reward selections tend to anchor optimal trajectories, enable recovery from deviations, and mitigate bottlenecks, with no single strategy dominating across domains; training-phase methods can approach near-optimal policies at substantially reduced labeling budgets. The findings demonstrate reward selection as a powerful paradigm for scaling RL in feedback-limited settings, with implications for RLHF, drug discovery, and other real-world problems where labeling costs are prohibitive.

Abstract

The ability of reinforcement learning algorithms to learn effective policies is determined by the rewards available during training. However, for practical problems, obtaining large quantities of reward labels is often infeasible due to computational or financial constraints, particularly when relying on human feedback. When reinforcement learning must proceed with limited feedback -- only a fraction of samples get rewards labeled -- a fundamental question arises: which samples should be labeled to maximize policy performance? We formalize this problem of reward selection for reinforcement learning from limited feedback (RLLF), introducing a new problem formulation that facilitates the study of strategies for selecting impactful rewards. Two types of selection strategies are investigated: (i) heuristics that rely on reward-free information such as state visitation and partial value functions, and (ii) strategies pre-trained using auxiliary evaluative feedback. We find that critical subsets of rewards are those that (1) guide the agent along optimal trajectories, and (2) support recovery toward near-optimal behavior after deviations. Effective selection methods yield near-optimal policies with significantly fewer reward labels than full supervision, establishing reward selection as a powerful paradigm for scaling reinforcement learning in feedback-limited settings.

Figures (6)

• Figure 1: Each row represents a data sample; shaded green rows indicate samples that have been labeled with rewards. The strategy $\mathcal{Q}_i$ determines which states to select for reward labeling. In the limited feedback setup, only a subset of states can be labeled. Different choices of reward-labeled subsets yield learnt policies of varying performances. The objective is to identify the subset that leads to the highest-performing policy.
• Figure 2: Problem setup for reward selection: The green arrows indicate the test phase, during which the reward selection strategy is evaluated. The blue arrows represent access to, and feedback from, an evaluator available within the training phase loop.
• Figure 3: Performance vs. training cost for selection strategies in Seaquest (60% feedback). Optimal strategies require prohibitive training cost (right), while cost-efficient and heuristic approaches trade off some performance (left). The dotted region indicates where cost-efficient strategies could emerge.
• Figure 4: Comparison of guided, visitation, and uniform heuristic selection strategies on four large-scale domains: Breakout, Freeway, Seaquest, and Asterix. For each domain, the plot shows the mean policy return with error bars indicating the standard error.
• Figure 5: Performance comparison of training-phase strategies and training-free guided on prototypical domains. Values show mean policy return over five test datasets (standard errors negligible). sequential-greedy achieves near-optimal performance, while guided is comparable to ES.