[Clarification request] Per-turn rewards in SKyRLGymGenerator agent_loop

[alex-dr]

I'm interested in multi-turn environments with turn-level rewards in the SkyRLGymGenerator.

From my reading of the code, only the final reward from the trajectory is used: https://github.com/NovaSky-AI/SkyRL/blob/main/skyrl-train/skyrl_train/generators/skyrl_gym_generator.py#L187 - the reward variable is updated with the most recent reward at each step, and the final value returned at the end of the method.

However, loss functions like GSPO explicitly support token-level advantages.

How am I supposed to give varying token-level rewards when only the final reward is ever returned, or is this not supported?

[tyler-griggs]

Hey @alex-dr, a great point! The trainer itself supports token-level rewards (link), and to test this code path I used a simple environment that maintained it's own token-level rewards. But, you're totally right, token-level rewards appear not to be supported by the SkyRLGymGenerator right now and we should update it.

We'd like to improve support for this! Can you share a little more about your setup/environment and we can help make it happen? For example, would this be helpful:

Update BaseTextEnvStepOutput to change the reward field from float to Optional[float]. On each step(), if the reward is non-None, we will append it to a list that racks each returned reward with which token id index it is associated with (ie, the final token id in the assistants response for that turn). Then, after the agent loop has completed, the SkyRLGymGenerator converts this to token-level rewards and passes to the trainer.

As described, this is basically a response-level reward. Do you want to assign even finer-grained rewards from the environment (e.g., token-level)?

[alex-dr]

The first setup you described sounds helpful. I do not want finer-grained reward assignment than that. I'm more interested in things like giving penalties for improperly formatted tool calls / rewards for successful ones, and possibly some NL verification/LLM feedback along the way.

However, directly broadcasting the reward to the tokens may not always be the appropriate behavior. Note for example that in the GRPO paper the process supervision version of the advantage calculation (4.1.3) tracks only a vector of rewards collected along the trajectory, and assigns per token advantage based on the sum of future rewards in the trajectory (after a group normalization). This may be different for other algorithms.

[tyler-griggs]

That makes sense! I believe if the SkyRLGymGenerator maintains this vector of (response-level) rewards along with the idx into the token_ids vector that each reward corresponds to, we should be able to support the process supervision described in the GRPO paper (and other similar custom advantage calculations based on response-level rewards).

Namely, SkyRL supports registering a custom advantage estimator in the trainer (code link). So, if the SkyRLGymGenerator passes back to the trainer a vector of token-level rewards where all rewards are 0 except for the token at the response boundaries, which hold the the response-level rewards (or "step-level rewards" as used in the GRPO paper section 4.1.3), then the registered custom advantage function can implement the process supervision method.

What do you think?

[alex-dr]

That makes sense to me. The alternative would be a list of (index, reward) tuples, but your structure seems more convenient to work with.

[tyler-griggs]

Okay great, I'll put this together

[tyler-griggs]

Hi @alex-dr, this change has been pushed in #226, so I will close this Issue. I hope it helps! Let us know if other challenges come up.

[alex-dr]

Thanks!