Reinforcement learning improves by splitting tasks into smaller parts
Reinforcement Learning with Decomposed Subtasks
Artificial IntelligenceMachine Learning
Summary
Some computer programs learn by trying things and getting a single overall score for success, which can make it hard to know what parts of their actions were helpful. The authors suggest breaking a big task into smaller subtasks and giving feedback on each subtask separately. This new way of learning helped programs do better on tasks where different skills matter and rewards come late. Their method works especially well for complicated tasks that require many steps and different abilities.
What this means in practice
- •For robotics teams: Train robots to learn complex tasks more efficiently by giving feedback on distinct skills within a long sequence of actions.
- •For software developers: Improve AI assistants that use multiple tools by breaking down long tasks into subtasks, making training more effective and interpretable.
Authors
Mattie Terzolo, Mikolaj Sacha, Ayan Sinha, Andrew Rabinovich
Abstract
Group Relative Policy Optimization (GRPO) and related policy-gradient methods for training language model agents collapse an entire multi-turn rollout into a single scalar trajectory reward before it enters the policy update. When the task composes distinct skills, especially under sparse and delayed environmental feedback, this collapsing is lossy: the optimizer must implicitly infer which competency drove the outcome and how that should change behavior. We argue the right primitive is not a better scalar but a decomposition: trajectory reward should be split along subtasks before it enters the policy update. We introduce Reinforcement Learning with Decomposed Subtasks (RLDS), whose core is Subtask-Decomposed Advantage Estimation (SDAE): a replacement for the scalar GRPO advantage that splits trajectory reward into per-subtask shares on a fixed taxonomy, computes a group-relative advantage per subtask, and distributes per-token credit by weighting each subtask's advantage by its importance, concentrating it around the step where a reflection marks that subtask's execution as consequential. We evaluate on four agentic benchmarks: FrozenLake (sparse grid navigation), HotpotQA (multi-hop QA, one retrieval tool), ScienceWorld (long-horizon embodied science), and DeepResearch (long-form research, four tools, composite rubric reward). Heterogeneity diagnostics emitted during training show where decomposition pays off - gains scale with subtask heterogeneity, largest on the high-heterogeneity tasks ScienceWorld (+11.5 points, paired-bootstrap 95% CI [+9.8, +13.3]) and FrozenLake (+9.8 points, [+7.0, +12.8]), and within noise on HotpotQA and DeepResearch, where the diagnostics predicted little to recover. ScienceWorld is also more compute-efficient under RLDS than scalar GRPO (-10.9% wall-clock per step), as long rollouts amortize the fixed reflect-and-grade overhead.