feat(reward): add configurable reward clipping and batch standardization for GSPO/GRPO

[Beisheng114]

Summary

This PR adds configurable reward clipping and per-batch standardization to AReno's GSPO/GRPO training pipeline, enabling operators to shape the reward distribution after scoring and before advantage computation. Three modes are supported — disabled (default, numeric identity), clip (fixed-range clamping), and standardize (per-batch z-score) — with raw and transformed reward distributions reported separately through logs, TensorBoard scalars, and CLI output.

Related Issue

#208

Quick Start

# Default behavior (disabled) — unchanged from existing GSPO/GRPO runs
areno train --ckpt Qwen/Qwen3-0.6B --dataset-path gsm8k:main \
  --reward-fn-path examples/math/math_verify_reward.py --algo gspo --tp-size 4

# Clamp noisy verifier rewards to a fixed range before GRPO advantages
areno train --ckpt Qwen/Qwen3-0.6B --dataset-path gsm8k:main \
  --reward-fn-path examples/math/math_verify_reward.py --algo grpo \
  --reward-transform-mode clip --reward-clip-min -1.0 --reward-clip-max 1.0 \
  --tp-size 4

# Standardize the batch reward distribution before GSPO advantages
areno train --ckpt Qwen/Qwen3-0.6B --dataset-path gsm8k:main \
  --reward-fn-path examples/math/math_verify_reward.py --algo gspo \
  --reward-transform-mode standardize --tp-size 4

What Changed

New Files

File	Lines	Description
areno/api/reward_transform.py	137	Core module: RewardTransformConfig + transform_rewards
tests/test_reward_transform_cpu.py	275	CPU test suite (21 test cases: unit + integration)

Modified Files

File	Changes
areno/api/trainer_config.py	Added reward_transform_mode, reward_clip_min, reward_clip_max, reward_transform_eps fields + reward_transform_config() factory on PolicyTrainerConfig
areno/api/models.py	Added optional TrainSequence.transformed_reward: float | None = None
areno/api/metrics.py	collect_train_batch_stats collects transformed_rewards only when set; record_training_stats emits rollout/transformed_reward_* scalars
areno/api/trainers/policy_only.py	New _transform_batch_rewards helper; both materialize methods apply transform before compute_group_advantages
areno/cli/train.py	4 CLI options (--reward-transform-mode/min/max/eps), _validate_reward_transform preflight, summary row, dashboard JSON fields
docs/cli/training.rst	Option docs, input contract, defaults, output fields, limitations, 2 copyable examples
docs/cli/observability.rst	rollout/transformed_reward_* TensorBoard tags + stage=reward_transform log line
tests/test_metrics_cpu.py	Transformed-reward metric separation test
tests/test_train_cli_config_cpu.py	11 CLI validation + summary tests; _options helper extended

Key Features

• Three modes — disabled (default, numeric identity), clip (clamp to [min, max]), standardize (per-batch z-score (r - mean) / (std + eps))
• Default backward compatible — disabled returns list(raw) exactly; no new TensorBoard scalars, log lines, or model fields appear
• Fail-fast validation — Invalid config rejected at CLI preflight before model/worker init; non-finite rewards raise ValueError with stage name + index (no sample payload exposed)
• Separate distributions — seq.reward stays raw; seq.transformed_reward set only when enabled; rollout/rewards_* vs rollout/transformed_reward_* kept distinct end-to-end
• Algorithm gating — GSPO/GRPO only; PPO/SFT/DPO rejected at validation (not silently ignored)
• Constant-batch guard — standardize with std=0 uses unit scale so output stays exactly zero and finite
• Empty-input guard — Early return skips numpy mean/std on zero-length arrays (no RuntimeWarning)
• PPC-safe — PPO overrides both materialize methods entirely; GAE reward path untouched

Reward Transform Flow

reward_fn scoring → _transform_batch_rewards (batch-wide) → compute_group_advantages (per-group)
                         ↑ disabled: identity               ↑ operates on transformed slices
                           clip: np.clip(raw, lo, hi)
                           standardize: (r-μ)/(σ+ε)

Stage	Raw Distribution	Transformed Distribution
seq.reward	Always raw	—
seq.transformed_reward	—	Only when enabled (None otherwise)
rewards_all (return)	Always raw	—
TensorBoard rollout/rewards_*	Always emitted	—
TensorBoard rollout/transformed_reward_*	—	Only when enabled (absent in disabled)
Log stage=reward_transform	raw[count mean std min max]	transformed[...]

Test Coverage

106 tests across 3 files, all CPU-only.

PYTHONPATH=. python -m pytest \
  tests/test_reward_transform_cpu.py \
  tests/test_train_cli_config_cpu.py \
  tests/test_metrics_cpu.py -v
# 106 passed, 1 warning in 1.86s

Test Class	Tests	Coverage
RewardTransformUnitTest	11	Disabled identity, clip clamps, standardize zero-mean/unit-std, constant→zeros, empty→[], empty-no-warnings, NaN raises with index, inf raises in disabled, extreme+clip finite, separate summaries
RewardTransformConfigTest	6	Unknown mode, clip requires both bounds, inverted bounds, non-finite bounds, non-positive eps, enabled property
PolicyOnlyMaterializeIntegrationTest	4	Disabled advantages match baseline, clip derives from clipped, standardize logs + sets transformed_reward, disabled emits no transform log
CLI config tests (reward transform)	11	Disabled by default, clip requires bounds, inverted bounds, bounds-without-mode, PPO/SFT rejected, clip shape, invalid eps (0/NaN), non-finite bounds, standardize summary, SFT n/a
MetricsUtilityTest	1	Raw and transformed rewards kept separate; disabled adds no transformed scalars

Review Fixes

During two review passes, 7 issues were found and fixed:

1. Empty-input numpy warnings — transform_rewards([], standardize) emitted Mean of empty slice RuntimeWarnings. Fixed: early-return guard before mode dispatch. Added test_empty_input_emits_no_numpy_warnings.
2. CLI eps validation gap — _validate_reward_transform checked eps but no CLI test exercised it. Added test_train_config_reward_transform_rejects_non_positive_eps + _rejects_nan_eps.
3. CLI non-finite bounds gap — Only tested at config level. Added test_train_config_reward_transform_clip_rejects_non_finite_bounds.
4. Standardize summary gap — Only clip mode tested in config summary. Added test_training_config_summary_shows_reward_transform_standardize.
5. Observability doc — observability.rst missing transformed_reward_* tags and stage=reward_transform log line. Updated.
6. Log-line assertions — Integration test didn't assert raw[ and transformed[ blocks in log. Added assertions.
7. GPU-validation docs — Issue requires documenting minimal GPU validation that remains. Added explicit docstring to integration test class.

Breaking Changes

None. Default reward_transform_mode="disabled" is a numeric identity. transformed_reward defaults to None, so no new TensorBoard scalars or log lines appear. seq.reward and rewards_all stay raw. All new config fields have safe defaults. Full CPU suite: 379 passed (1 unrelated failure from missing openai package).

Files Modified

areno/api/
├── reward_transform.py              [NEW] 137 lines — RewardTransformConfig + transform_rewards
├── trainer_config.py                [MODIFIED] 4 new fields + reward_transform_config() factory
├── models.py                        [MODIFIED] TrainSequence.transformed_reward field
├── metrics.py                       [MODIFIED] collect + record transformed_reward distribution
└── trainers/
    └── policy_only.py               [MODIFIED] _transform_batch_rewards + two-pass materialize
areno/cli/
└── train.py                         [MODIFIED] 4 options + _validate_reward_transform + summary + JSON
docs/cli/
├── training.rst                     [MODIFIED] option docs, contract, limitations, 2 examples
└── observability.rst                [MODIFIED] transformed_reward tags + reward_transform log line
tests/
├── test_reward_transform_cpu.py     [NEW] 275 lines, 21 test cases
├── test_metrics_cpu.py              [MODIFIED] transformed-reward metric separation test
└── test_train_cli_config_cpu.py     [MODIFIED] 11 CLI validation + summary tests

Hardware Limitations

• CPU tests: All 106 tests run without a GPU (verified on Python 3.11.15 with CPU-only torch)
• GPU/distributed: Per-rank statistics semantics (each TP rank shapes its own reward shard) documented as a limitation; orchestration logic faked in integration test via object.__new__(PolicyOnlyTrainer) with SimpleNamespace config
• GPU validation remaining: End-to-end GSPO/GRPO training with real model/rollout/backend to confirm transformed rewards flow through the loss without shape/dtype mismatches — requires CUDA hardware