fix(ppo): reject implicit partial group normalization

[EazyReal]

Description

Current main already includes variable-size group normalization through #1454. This PR keeps the remaining invariant: when group-level Normalization is called without explicit group_sizes, it is using fixed-size groups, so the batch size must be divisible by group_size.

Before this guard, a direct fixed-stride call with a partial tail could leave tail rows with zero-filled normalization state. The PPO rollout path now passes TrajBatchMeta.traj_group_sizes, so this is a small boundary hardening for direct callers and future refactors.

Superseded Context

The original broader #1415 head was 3e3ab94b257acf1a77dc5a9fdecf7cfea2a6f7f1. That broader variable-group normalization change was superseded by #1454, merged to main as bbc10f0e95bf350d71169efcb2facb9756f1238b (PR commit 375cc12a4b91da50614aa9d2211ab9995475c942). This PR is now only the remaining non-duplicative guard.

Validation

• uv run pytest -q tests/test_reward_norm_variable_group.py tests/test_adv_norm_config.py (214 passed)
• uv run ruff check areal/utils/data.py tests/test_reward_norm_variable_group.py
• uv run ruff format --check areal/utils/data.py tests/test_reward_norm_variable_group.py
• git diff --check origin/main...HEAD

Related Issue

Part of #1419.

[gemini-code-assist]

Code Review

This pull request introduces support for group-level normalization with unequal or partial groups by passing actual per-group row counts (group_sizes) through the PPO actor and normalization utilities. It replaces batched_call with explicit batch concatenation and splitting to recover trajectory group sizes, updates the normalization logic to resolve dynamic group boundaries, and adds comprehensive unit tests. The review feedback suggests adding a guard for empty input data in compute_advantages to prevent potential crashes, and validating that all elements in group_sizes are strictly positive to avoid runtime errors during tensor expansion.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[EazyReal]

Thanks for the review. Addressed:

• Validate group_sizes are positive — done: _resolve_group_bounds now raises if any element is < 1, alongside the existing sum check.

Respectfully holding off on:

• Empty-data guard in compute_advantages — this is pre-existing behavior (the previous batched_call(self._compute_advantages, data) had the identical empty-input path), it isn't reachable in the training loop (prepare_batch blocks until a full batch of accepted results), and the same assumption applies to compute_logp/compute_values. Guarding only this method would be inconsistent; if empty-batch handling is wanted it belongs in a separate, deliberate change. Happy to add it if you'd prefer.

[EazyReal]

@garrett4wade @sitabulaixizawaluduo @rchardx @fishcrap Could I get a quick review on this partial-group PPO normalization fix? It uses actual group sizes for under-filled rollout groups.

[EazyReal]

Retargeted this to the remaining non-duplicative guard on current main. @sitabulaixizawaluduo if you have bandwidth, would appreciate a review on the fixed-stride group-normalization boundary check.

[github-actions]

This pull request has been automatically marked as stale because it has not had recent activity within the last 14 days.

Please add a comment or push new commits to keep it active.

Thank you for your contribution!

[EazyReal]

The hazard this guards is still live on main (2ee4904): with group_sizes=None and a batch not divisible by group_size, Normalization._build_group_slices covers only bs // group_size groups and __call__ returns the uncovered tail as x/(0+eps) — reproduced on main: bs=3, group_size=2 silently gives the orphan row an advantage of ~2e5. #1454 added the explicit group_sizes branch but left this fallback unguarded. This PR is the minimal fix — a 5-line divisibility guard plus a regression test, mirroring the invariant train_controller.py already enforces one layer up. One note: open #1563 currently carries the same hunk; whichever lands first, the other drops it. @sitabulaixizawaluduo could you take a look?