Skip to content

feat: reward shaping layer with success-gated length penalty - #3376

Draft
hallerite wants to merge 1 commit into
mainfrom
feat/reward-shaping
Draft

feat: reward shaping layer with success-gated length penalty#3376
hallerite wants to merge 1 commit into
mainfrom
feat/reward-shaping

Conversation

@hallerite

Copy link
Copy Markdown
Member

Summary

Pulls the length penalty out of GRPOAlgorithm into a reward-shaping layer that runs before credit assignment, and gates it on success.

Why

The penalty modifies the reward; the algorithm decides what to do with the reward afterwards. Hacked into GRPO it was unavailable to max_rl (a one-line variant of GRPO), rae and hierarchical_grpo. And it penalized every rollout in a group — a wrong long answer paid the same as a right long answer, so among failures shorter-wrong out-ranked longer-wrong: "if you're going to fail, fail fast", the opposite of what you want on hard problems.

What

  • [[orchestrator.reward_shaping]] — a list of shapers (type = "length_penalty" is the first), inheriting to each [[orchestrator.train.source]] like algo does (unset inherits, [] disables). GRPOAlgoConfig.length_penalty is removed.
  • RewardShaper.shape_group (orchestrator/reward_shaping.py) runs in TrainSink.process_group right before algorithm.finalize_group. Each shaper records an additive term on every trainable trace in trace.reward_shaping[<type>] (new verifiers field, companion PR below).
  • training_reward(trace) = trace.reward + sum(trace.reward_shaping.values()) is what grpo, echo, max_rl, rae and hierarchical_grpo now score against. trace.reward stays the env's verdict, so solved_*, pass@k gating, curriculum pools and evals are unaffected.
  • Success gate: penalty_i = 1[r_i ≥ success_threshold] · mean(r) · frac_i, success_threshold = 1.0 by default (lower it for partial-credit envs). mean(r) stays as the difficulty scaler.
  • Each shaper's term is logged under reward_shaping/<type>.
  • The five in-tree configs using the penalty (configs/debug/concurrency.toml, examples/advanced/{glm-4.5-air/*,nemotron-3-super/swe}.toml) are migrated with the same weights — note they now gate on success by default.
  • Docs: new Reward Shaping section in docs/algorithms.md; docs/training.md and skills/configs updated.

Companion

Depends on PrimeIntellect-ai/verifiers#2436 (Trace.reward_shaping); the deps/verifiers submodule is bumped to that branch's commit (based on the currently pinned d4f3d5a8c, so no other verifiers changes ride along).

Tests

  • tests/unit/orchestrator: 87 passed. Length-penalty tests now drive LengthPenalty.shape_groupGRPOAlgorithm; one new test covers the success gate and threshold override.
  • test_env_algo_overrides_top_level extended to cover the reward_shaping inherit / [] rule.

🤖 Generated with Claude Code

Pull the length penalty out of GRPOAlgorithm into a RewardShaper layer that
runs before credit assignment, gate it on success (success_threshold, default
1.0), and score every reward-consuming algorithm against training_reward.
Bump deps/verifiers to the Trace.reward_shaping companion commit.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@hallerite
hallerite force-pushed the feat/reward-shaping branch from e09bce1 to d1c84bc Compare August 24, 2026 23:13
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant