RLDS: Rethinking Credit Assignment in Reinforcement Learning via Subtask Decomposition

The Bottleneck of GRPO: What Gets Lost in a Scalar Reward

When training language model agents, GRPO (Group Relative Policy Optimization) and related policy gradient methods share a common approach: compressing an entire multi-turn interaction trajectory into a single scalar reward before feeding it into the policy update. This seemingly reasonable simplification reveals a clear weakness when handling tasks that require composite skills.

When a task combines multiple distinct capabilities — especially under sparse or delayed environment feedback — this "collapse" is lossy. The optimizer is forced to implicitly infer: which capability actually drove the final outcome, and how should it be adjusted? In other words, a single score can't tell the model "what went right and what went wrong," leaving credit assignment fundamentally ambiguous.

The core thesis of this arXiv paper is: the right solution isn't to design a "better scalar" — it's to decompose. Trajectory rewards should be broken down along subtask dimensions before entering the policy update.

GRPO (Group Relative Policy Optimization) is a policy gradient algorithm for LLM alignment and reinforcement learning fine-tuning, introduced by the DeepSeek^Ad team in 2024 and widely applied in reasoning model training (e.g., DeepSeek-R1). Its core idea: sample a group of outputs for the same input, use the group's average reward as a baseline, and compute each output's advantage relative to that baseline — avoiding the need to train a separate critic network. This design makes GRPO lighter than PPO and particularly suited for generative tasks. However, it presupposes that the reward signal is sufficiently "informative" — that a single scalar can represent the quality of an entire trajectory. When a task involves multiple heterogeneous sub-skills, this assumption causes the optimizer to lose the ability to distinguish "which sub-skill caused the overall failure," amplifying the credit assignment problem.

The Core Mechanism of RLDS and SDAE

The paper introduces RLDS (Reinforcement Learning with Decomposed Subtasks), whose core component is Subtask-Decomposed Advantage Estimation (SDAE). It is designed as a drop-in replacement for the scalar GRPO advantage function, and its workflow breaks down into three steps.

Splitting Rewards by a Fixed Taxonomy

SDAE first distributes the trajectory reward across individual subtasks according to a fixed subtask taxonomy, producing a reward share for each subtask. This expands the feedback that was previously collapsed into a single number back into a set of semantically meaningful components.

Computing Group-Relative Advantage per Subtask

After splitting, the method computes a group-relative advantage for each subtask independently. This preserves GRPO's approach of using within-group samples as a baseline, but refines the comparison granularity to the subtask level rather than the full trajectory.

Assigning Token-Level Credit by Importance

Finally, SDAE assigns per-token credit by weighting each subtask's advantage according to its importance. Credit is concentrated around the "reflection" token marking the step where that subtask's execution had a critical influence. This directs gradient signals more precisely to the decision points that actually matter, rather than spreading them uniformly across the entire sequence.

Per-token credit assignment is a classic challenge in sequential decision-making RL. In the context of language model training, a single response may contain hundreds of tokens, yet only a handful of key decision points truly determine the final reward. Spreading the reward signal uniformly across every token means large numbers of outcome-irrelevant tokens also receive gradient updates, creating "noisy gradients." SDAE uses a "reflection" token to locate which step had a critical influence on a given subtask's execution, concentrating gradients at those positions. This is effectively a lightweight form of temporal credit assignment — conceptually analogous to the role of an advantage function in RL, but refined to the cross-dimension of subtasks and tokens.

Experimental Results Across Four Benchmarks

The authors evaluated RLDS on four representative agent benchmarks spanning different task structures — from sparse navigation to long-horizon scientific research:

• FrozenLake: Sparse-reward grid navigation
• HotpotQA: Multi-hop question answering with a single retrieval tool
• ScienceWorld: Long-horizon embodied science tasks
• DeepResearch: Long-form research involving four tools with a composite rubric reward

One particularly notable experimental design choice is the heterogeneity diagnostics logged during training, which reveal when decomposition is genuinely valuable. The conclusion is clear: gains scale with subtask heterogeneity.

Improvements are most pronounced on high-heterogeneity tasks — +11.5 points on ScienceWorld (paired bootstrap 95% CI: [+9.8, +13.3]) and +9.8 points on FrozenLake ([+7.0, +12.8]). On HotpotQA and DeepResearch, improvements fall within the noise range — which is exactly consistent with the diagnostics: these two tasks simply don't have much recoverable information to begin with.

This "diagnose first, validate with results" presentation makes the method's applicability interpretable rather than asserting universal effectiveness.

These four benchmarks each represent structurally distinct tasks. FrozenLake is a classic sparse-reward grid world where the agent must navigate to a goal with almost no intermediate feedback. HotpotQA requires multi-hop reasoning across multiple documents using a single retrieval tool, with relatively low subtask heterogeneity. ScienceWorld simulates long-horizon science experiments in an embodied environment, requiring multi-step planning and execution across heterogeneous skills including perception, reasoning, and action. DeepResearch mimics real-world literature review workflows using four tools (search, summarization, citation, etc.) and evaluates report quality with a composite rubric reward. This diversity ensures the experimental conclusions are reliable across varying feedback densities and skill compositions.

An Additional Gain: Computational Efficiency

Beyond task performance, RLDS also demonstrates better computational efficiency than scalar GRPO on ScienceWorld — reducing wall-clock time per step by 10.9%. The reason: longer trajectories amortize the fixed overhead of the "reflect-and-grade" step more effectively. The longer the trajectory, the more thoroughly this additional cost gets spread out.

This is practically meaningful. One might naturally worry that finer-grained credit assignment would introduce additional computational burden — but in long-horizon task settings, RLDS actually comes out ahead on efficiency.

Significance and Boundaries of the Method

RLDS contributes more than just improved scores on two benchmarks — it fundamentally reconsiders the basic unit of reward processing in policy gradient methods. Shifting the optimization direction from "design a better scalar" to "how to decompose the reward" represents a conceptual shift.

At the same time, the paper maintains a measured stance on the method's scope: the gains from decomposition strongly depend on the subtask heterogeneity of the task. For tasks with homogeneous skill requirements and relatively dense feedback, fine-grained decomposition offers limited value. This honest delineation of boundaries actually strengthens the credibility of the conclusions.

For researchers focused on LLM agent training, RLDS offers a transferable paradigm: before investing effort in improving reward scalars, ask first — is this task's reward worth decomposing?