Skip to content

refactor(algorithms): extract advantage estimators - #27

Draft
leviking98z-rgb wants to merge 1 commit into
refactor/rollout-capability-interfacesfrom
refactor/advantage-estimator-components
Draft

refactor(algorithms): extract advantage estimators#27
leviking98z-rgb wants to merge 1 commit into
refactor/rollout-capability-interfacesfrom
refactor/advantage-estimator-components

Conversation

@leviking98z-rgb

Copy link
Copy Markdown
Owner

Summary

Extracts reward-to-advantage policy from Part and trainer-specific helpers into a reusable algorithm-layer component.

  • Adds typed AdvantageBatch, AdvantageEstimate, and structural AdvantageEstimator contracts.
  • Provides parity-preserving grouped/global GRPO normalization and the existing 1D/2D GAE implementation as components.
  • Migrates synchronous/async AR, diffusion, PE, unified-model, barrier/partial/async agentic trainers, with an optional root advantage: Hydra component.
  • Removes Part.compute_advantages, the standalone normalizer module, and the old types-owned GAE module.
  • Extends the framework guard so advantage policy cannot silently return to the wire types.

This removes 421 lines of old helpers and trainer policy. The reusable implementation, typed contracts, compatibility wiring, docs, and persistent guard make the slice net +125 lines; its purpose is to stop GDPO, PPO/value-head, and future estimators from adding branches to every trainer.

Related Issue

N/A

Test Plan

  • SKIP=no-commit-to-branch pre-commit run --all-files --show-diff-on-failure — passed, including Ruff, recipe targets, experimental boundaries, architecture boundaries, and framework contracts.
  • python -m compileall -q unirl lint/check_framework_contracts.py — passed.
  • python lint/check_architecture_boundaries.py — passed: 473 files / 4,264 internal import edges satisfy 8 boundaries.
  • python lint/check_framework_contracts.py — passed: 7 rollout engines, 6 weight syncs, 19 model pipelines, 2 train backends, 2 wire types, 5 advantage contracts, 4 loop programs, and 13 entrypoints conform.
  • python lint/check_recipe_targets.py — passed: 2,311 recipe _target_ paths resolve.
  • One-off NumPy-backed torch-stub parity harness (not committed) — passed 9 cases covering grouped population variance, global sample std, grouped means/global std, mean-centering only, agentic non-finite exclusion, interleaved completion groups, 1D GAE, returns, and 2D masked GAE; it also rejects a None group id.
  • GPU rollout/training smoke — not run; this environment has no torch/Ray/Hydra stack or accelerator allocation.

Compatibility / Risk

  • Existing recipes preserve their current defaults through adv_normalization_scope, normalize_adv_by_std, and adv_use_global_std. An explicit advantage: block takes precedence.
  • No checkpoint or Sample/Part wire-field migration is required; Part still carries rewards, component rewards, and resulting advantages.
  • This intentionally removes the Python APIs Part.compute_advantages, unirl.algorithms.normalizers, and unirl.types.advantages. Out-of-tree callers should use estimate_part_advantages or an AdvantageEstimator.
  • Agentic defaults preserve non-finite failed-trajectory exclusion. A custom agentic estimator must define equivalent non-finite handling.
  • GeneralizedAdvantageEstimator is the algorithm contract only; a value-based trainer must supply values and masks.

Reviewer Notes

Checklist

  • I reviewed the changed code and removed unrelated/generated artifacts.
  • I updated tests, docs, and configs where needed, or explained why not.

@github-actions github-actions Bot added the wip label Jul 31, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant