[fix][train][step-wise] Broadcast step-wise advantage with each step's own response_mask

[CharlieFRuan]

In step-wise training, advantages = last_step_advantages[traj_ids] copied the last step's already-masked tensor verbatim to every step of the trajectory. Because each step has its own right-aligned response_mask at its own token positions (prompts and responses vary in length across turns), the scalar advantage landed at the last step's token positions — which almost never overlap with the earlier step's own response positions. Earlier steps silently received advantages = 0 at their loss_mask=1 positions, so roughly (N-1)/N of step-samples are trained incorrectly.

Fix:

• Call the estimator with an all-ones mask so it returns the scalar advantage at every position (outcome-based estimators produce scalar * mask).
• Broadcast by indexing into traj_ids, then multiply by each step's own response_mask so values land at the right token positions.
• Reject step_wise_trajectories=True + gae/reinforce++ in validate_cfg: temporal estimators produce per-token advantages that the scalar-broadcast model can't represent, and the step-wise path only sees the last step's slice anyway (no cross-step credit).
• Add regression test test_calc_advantages_and_returns_step_wise_broadcast with a realistic right-aligned batch, plus a parametrized test_validate_cfg_step_wise_estimator_compatibility covering all 5 built-in estimators.
• Document the outcome-only constraint in the step-wise tutorial.

Fixes #1492.

Testing

https://wandb.ai/sky-posttraining-uc-berkeley/skyrl-search-padding/reports/PR-1285-search-r1-PAD-left-only-baseline-after-PR-and-step-wise---VmlldzoxNjE3MTA0OQ

Since we will run into #1501 at HEAD, we make this hacky hot fix, mainly the pad_batch() change: CharlieFRuan@f2783a4

    if self.cfg.generator.step_wise_trajectories:
        n_samples = self.cfg.generator.n_samples_per_prompt
        pad_target = max(
            self.cfg.trainer.policy_mini_batch_size * n_samples,
            self.cfg.trainer.critic_mini_batch_size * n_samples,
            dp_size,
        )
    else:
        pad_target = dp_size
    pad_size = math.ceil(training_input.batch_size / pad_target) * pad_target - training_input.batch_size

We compare against the red curve from #1285

Note that the actual curve should converge slower than this after we merge in #1483, in which each mini batch will truly be one optimizer step. Currently we can have multiple optimizer step due to the number of turns can inflate the number of sequences in each mini batch.

The reward curve being higher is likely due to:

1. The fix we apply in this PR
2. In the previous PR (red line), we silently strip the amount of data and only train num_mini_batches = len(data) // mini_batch_size, which will no longer be an issue with the pad_target hacky fix above, or properly fixed after [fix][train] Prompt-based mini-batching for step-wise training #1483

---

[gemini-code-assist]

Code Review

This pull request fixes an issue with advantage broadcasting in step-wise training by ensuring trajectory-level rewards are correctly distributed across all steps using per-step response masks. It also adds configuration validation to restrict step-wise training to outcome-based advantage estimators (GRPO, RLOO, MAXRL), updates the documentation, and includes new regression tests. I have no feedback to provide.

[devin-ai-integration]

✅ Devin Review: No Issues Found

Devin Review analyzed this PR and found no potential bugs to report.

View in Devin Review to see 4 additional findings.

[SumanthRH]

Thanks! Would be good to get a E2E run!