[Fix] Fix pass_at_k missing for SkRLGymGenerator.agent_loop flow due to token-level rewards

[CharlieFRuan]

In #226, we started building per-token reward for the agent_loop() codepath to enable per-step reward. However, in get_metrics_from_generator_output(), we do not compute pass_at_n for token-level rewards:

def get_metrics_from_generator_output(
    generator_output: GeneratorOutput, uids: List[str]
) -> Tuple[float, Optional[float]]:
    ...
    if isinstance(rewards[0], list):
        # We just compute mean over sequence reward.
        # TODO: We should make metrics customizable by the environment
        mean_raw_reward = float(np.mean([sum(seq_rewards) for seq_rewards in rewards]))
        pass_at_n = None  # not computed for token-level rewards since it's ill-defined
    else:
        ...

This PR resolves this by still using per-trajectory reward when all intermediate rewards are None.

[gemini-code-assist]

Code Review

This pull request addresses an issue where pass_at_k metrics were not being calculated for the agent_loop flow due to token-level rewards. The introduced logic correctly distinguishes between per-trajectory and per-token reward scenarios, ensuring that rewards are formatted as a float when appropriate, which allows for the pass_at_k calculation. The implementation is sound and effectively resolves the problem. I've included one suggestion to refactor a loop into a more concise and readable single-line expression.

[CharlieFRuan]

/gemini review

[gemini-code-assist]

Code Review

This pull request effectively addresses the issue of pass_at_k metrics not being computed for certain reward structures. The change in skyrl_train/generators/skyrl_gym_generator.py correctly distinguishes between per-trajectory and per-token rewards by checking if intermediate rewards are None. This ensures that per-trajectory rewards are passed as a single float, enabling the metric calculation. The updated tests in skyrl-train/tests/cpu/generators/test_skyrl_gym_generator.py are well-designed, using parameterization to validate both scenarios and confirming the fix. The code is clear, correct, and well-tested.