Test Difficulty Driven Reward

• Test Difficulty Driven Reward is a reinforcement learning approach that assigns graded rewards based on test difficulty groups to overcome sparse reward challenges.
• It employs both strict and soft reward schemes to incentivize progressive mastery of increasingly difficult programmatic tests in LLM-based code generation.
• Empirical results show that TDDR improves sample efficiency and pass rates on benchmarks compared to traditional binary reward methods.

Test Difficulty Driven Reward (TDDR) is a reinforcement learning strategy designed to address reward sparsity and stagnation in program synthesis and algorithmic reasoning tasks, particularly for LLMs fine-tuned via reinforcement learning over code. Rather than issuing binary rewards only when all test cases pass, TDDR introduces a graded, instance-specific reward based on the difficulty of each test case, determined by pass rates across a pool of models. This produces denser, curriculum-style reinforcement signals that encourage both exploration and policy improvement on hard-to-solve programmatic benchmarks (Xiaomi et al., 12 May 2025).

1. Motivation and Background

Sparse reward signals are a primary bottleneck in reinforcement learning for code generation and mathematical reasoning. Standard evaluation assigns reward 1 only if all test cases (T) are passed, and 0 otherwise. This leads to vanishing gradients on difficult problems, as most rollouts receive no reward signal. Moreover, naively rewarding partial test case success can be misleading if all test cases have the same weight, as easier cases may dominate and dilute policy updates.

Test Difficulty Driven Reward was introduced in the MiMo-7B project to alleviate these problems within group-based RL (Group Relative Policy Optimization, GRPO) for mathematics and code tasks. By incentivizing incremental progress on harder parts of the test suite and dense feedback even for partial solutions, TDDR fosters sample efficiency and stable policy learning in challenging domains (Xiaomi et al., 12 May 2025).

2. Construction of Difficulty-Aware Test Groups

In the TDDR scheme, each code problem is associated with T discrete unit tests. These tests are clustered into G groups of increasing difficulty based on aggregated per-test pass rates across multiple models (lower pass rates correspond to higher difficulty). Groupings can be constructed by hierarchical clustering or percentile-based binning of test success frequencies observed in model rollouts.

For each group $g\in\{1,\ldots,G\}$, a group-specific sub-score $s_g$ is assigned, typically using an IOI-style allocation. Each group thus represents a contiguous difficulty level, and each test $t$ belongs to some group $g(t)$. Grouping transforms the reward landscape to reflect not only partial correctness but challenge-specific attainment.

3. Reward Formulation

Two scoring schemes are implemented:

• Strict TDDR: Reward increments only when all tests in levels $\leq g$ have passed.

$$r_{\mathrm{strict}} = \sum_{g=1}^G s_g \cdot \mathbf{1}[\text{all tests in levels} \leq g\, \text{passed}]$$

This approach encourages models to prioritize and master easier test cases before progressing to harder ones.

• Soft TDDR: Each test $t$ contributes independently, weighted according to the difficulty group:

$$r_{\mathrm{soft}} = \sum_{t=1}^{T} w(d_t) \cdot \mathbf{1}[\text{test}\; t\; \text{passed}]$$

where $w(d_t) = s_g / | \text{group}_g |$ for $t \in \text{group}_g$. This delivers a denser reward, particularly beneficial for very challenging problems with large T.

Both schemes provide partial credit on complex instances that would otherwise yield no positive feedback under traditional all-or-nothing criteria. Empirical evidence demonstrates that this dense, difficulty-stratified reward structure alleviates reward sparsity and improves both the efficiency and ultimate performance of RL-tuned LLMs on code and math (Xiaomi et al., 12 May 2025).

4. Integration with Reinforcement Learning Pipelines

TDDR is natively integrated into the GRPO framework in MiMo-7B-RL's post-training phase. Rollouts for each problem generate multiple candidate solutions, each receiving a shaped reward according to the problem-specific TDDR formula. The normalized, advantage-weighted rewards are then used to update the model policy via the GRPO objective:

$s_g$0

with $s_g$1 being the normalized advantage derived from TDDR-based rewards.

Key stabilizing techniques include:

• Dynamic exclusion of problems with passrate 0 or 1 from mini-batches (to keep all examples on-policy and gradient-contributing)
• "Easy-data" re-sampling from a buffer of solved tasks at controlled rates, ensuring coverage and preventing policy collapse late in training (Xiaomi et al., 12 May 2025).

5. Empirical Results and Efficacy

The application of TDDR in MiMo-7B-RL yielded substantial gains across multiple code and math benchmarks. On LiveCodeBench v6, MiMo-7B-RL achieved 49.3% pass@1 compared to 46.8% for OpenAI o1-mini and 39.1% for QwQ-32B, reflecting state-of-the-art performance for open-source models of this scale. Importantly, the dense TDDR mechanism allowed models to make incremental progress on previously unsolved or sparse-reward domains, surpassing even much larger (32B parameter) baselines (Xiaomi et al., 12 May 2025).

6. Broader Context, Extensibility, and Limitations

TDDR exemplifies a class of reward shaping methodologies that use population-level or roll-out-level statistics to adapt reinforcement learning signals to the problem landscape. Such methods are most naturally applied in domains with a natural test suite structure, such as program synthesis, algorithmic reasoning, and possibly in selected vision-language or robotics settings provided deterministic subtask evaluation.

A plausible implication is that similar dense and difficulty-aware reward routing could accelerate RLHF and similar paradigms for multimodal, math, and theorem-proving settings if appropriate per-unit evaluation signals are available.

A limitation is that TDDR depends on the existence of diverse and graded test cases for each problem. Its utility is limited in tasks where performance is truly atomic or where no meaningful test stratification can be computed. Additionally, TDDR can introduce reward hacking risks if test clustering is not robust to model-induced artifacts.

7. Summary Table: Key Features of Test Difficulty Driven Reward

Aspect	Traditional Reward	Test Difficulty Driven Reward
Positive reward issued	All tests passed	Partial credit per test/group
Reward granularity	Binary (0/1)	Dense, stratified by difficulty
Curriculum structure	Absent	Emergent via difficulty grouping
RL sample efficiency	Low on hard tasks	High, due to density
Applicability	Code, math tasks	Code, math, or any testable suite

The TDDR formulation and deployment in the MiMo-7B RL pipeline constitute a pivotal advancement in reward engineering for LLM-based program synthesis, offering a reproducible mechanism for mitigating reward sparsity and facilitating curriculum-aligned exploration in challenging code and reasoning domains (Xiaomi et al., 12 May 2025).