[SkyRL-Gym] Make SQL and Search envs return None for intermediate steps

[erictang000]

Overview

After #226, the SkyRLGymGenerator expects turn level rewards to be None if the env uses trajectory level rewards. After #271, this causes issues computing metrics for SQL and Search envs. Setting intermediate reward to None and making it an optional type fixes this.

Additionally, handles case in SkyRLGymGenerator where we exit the agent loop due to hitting max length, which results in rewards being all None - this is fixed by setting the last reward to 0.

[CharlieFRuan]

/gemini review

[gemini-code-assist]

Code Review

This pull request correctly modifies the SQL and Search environments to return None for intermediate step rewards, aligning with the SkyRLGymGenerator's expectations. The change to handle episode truncation in SkyRLGymGenerator by setting a None reward to 0.0 is also a good addition. My review includes a few suggestions to correct the return type hints in the modified environment methods, which were missed, and a minor documentation consistency fix.

[gemini-code-assist]

With this change, the function can now return None. Please update the function's return type hint on line 45 from float to Optional[float] to reflect this. Optional is already imported in this file.

[gemini-code-assist]

This function now returns None for intermediate steps. The return type hint on line 78 should be updated from float to Optional[float] to match this change. You'll need to add Optional to your imports from the typing module (e.g., from typing import Optional).

[gemini-code-assist]

For consistency with Python's type hinting syntax, it would be clearer to use Optional[float] instead of Optional[Float].

Suggested change

	- reward: Optional[Float] reward for the action, None if intermediate steps have no reward
	- reward: Optional[float] reward for the action, None if intermediate steps have no reward

[CharlieFRuan]

As discussed offline, will revert #271 first to avoid the error that Eric is observing from happening.

The error is due to this line:

SkyRL/skyrl-train/skyrl_train/generators/utils.py

Line 65 in 2e7aba9

mean_raw_reward = float(np.mean([sum(seq_rewards) for seq_rewards in rewards]))

Where it was that sum(seq_rewards) doesn't work bc seq_rewards was a float, but rewards[0] is a list, which means some were lists and some were floats. This is because intermediate reward for sql and search are 0, so some trajectory's reward is [0, 1], and some trajectory's reward can be [1]. The first goes into else, the second goes into if:

SkyRL/skyrl-train/skyrl_train/generators/skyrl_gym_generator.py

Lines 292 to 298 in 2e7aba9

	if all(reward is None for reward, _ in per_step_rewards[:-1]):
	# If all rewards besides the last one are None (i.e. per-trajectory reward), we keep it as a float
	reward_out = per_step_rewards[-1][0]
	else:
	# Otherwise build token-level rewards placed at assistant turn boundaries
	token_level_rewards: List[float] = [0.0] * len(response_ids)
	for step_reward, idx in per_step_rewards:

Reverting #271 will prevent the error at the cost of not seeing pass_at_n metric. Which we will follow up with a fix.

[erictang000]

closing and tracking actual fix in #311