Lego-RL

Difficulty Filtering

Measure per-instance difficulty by rollout, then select the band that carries gradient

Difficulty is measured, not assumed. The stage runs the agent over the pool with the policy frozen, records how many attempts solved each instance, and keeps the band that produces gradient.

  • Input — a task index, and a policy to measure it with.
  • Output — an index whose rows have a pass rate strictly between 0 and 1.

Measure by rollout

Run batch inference over the pool with N_TRIALS attempts per instance at training temperature. The run writes per-instance counts to <RESULTS_DIR>/summary.csv and the default band to OUTPUT_INDEX.

The distribution is the reason this stage exists. Of the 14,415-instance OpenSWE shard, only 760 were solved at least once in 4 attempts. Under group-relative advantages, an instance solved 0 or 4 times out of 4 produces a zero-variance group and therefore no gradient at all. Roughly 90% of that shard would have contributed nothing but wall-clock time.

Select the band

OUTPUT_INDEX already keeps every instance with 0 < pass_rate < 1, dropping the all-pass and all-fail ends in one move.

For a narrower band, read the counts and join them to the pool on instance_id:

import pandas as pd

pool    = pd.read_parquet("train_openswe_all_qwen_merged_clean_2699.parquet")
summary = pd.read_csv("results/node0/summary.csv")
# columns: instance_id, n_completed, n_pass, n_fail, pass_rate, rewards

keep = set(summary.loc[summary.n_pass.between(1, 2), "instance_id"])
band = pool[pool["extra_info"].map(lambda e: e["instance_id"] in keep)]

band.to_parquet("train_openswe_2699_qwen36_27b_r1to2_951.parquet")

Keeping difficulty out of the index itself is deliberate. Counts are per model and per sampling temperature, so the same pool yields a different band for every policy you measure it with, while the index stays a stable pointer table.

Point a training config's TRAIN_FILES at the result. Because band selection is the last thing that happens, comparing bands is a controlled experiment: the difficulty ablation shipped in scripts/train/configs/ablation951_*.env holds instance count fixed across arms and varies only the band (solved 1–2, solved 2–3, and a random draw from the full pool as the baseline).

Draw the baseline arm from the whole pool, not from the instances that carry signal. A random draw from the clean pool is roughly 10% solvable, and that is exactly the number a difficulty filter has to beat.

On this page