fix(examples/grpo): build rollout_infos before _close_and_remove in sudoku scheduler (#9745)

[Anai-Guo]

What

Fixes the first defect reported in #9745: in the OpenEnv Sudoku GRPO example, the parse-failure branch of SudokuScheduler.on_turn_end builds its rollout_infos dict after calling _close_and_remove(uuid).

OpenEnvScheduler._close_and_remove (in swift/rollout/multi_turn.py) pops self._total_rewards[uuid] and self._step_rewards[uuid]. So when the model emits an unparseable action:

self._total_rewards[uuid] = self._total_rewards.get(uuid, 0.0) - 1.0
self._step_rewards.setdefault(uuid, []).append(-1.0)
await self._close_and_remove(uuid)          # pops _total_rewards[uuid] / _step_rewards[uuid]
return {
    'done': True,
    'rollout_infos': {
        'total_reward': self._total_rewards[uuid],        # KeyError
        'step_rewards': list(self._step_rewards.get(uuid, [])),  # always []
        ...

self._total_rewards[uuid] raises KeyError (and step_rewards silently comes back empty). Since unparseable model output is common early in GRPO training, this crashes rollouts.

Fix

Snapshot rollout_infos before cleanup — mirroring the already-correct terminal path at the end of on_turn_end (which builds rollout_infos first, then calls _close_and_remove). Minimal, behavior-preserving for the happy path.

Note on the second reported issue

#9745 also reports an OpenEnv session leak when a trajectory stops because it hit --max_turns (or finish_reason == 'length') rather than the env returning done=True. That one is a framework-level gap: MultiTurnScheduler.run() returns as soon as should_stop is set and never invokes a post-trajectory cleanup hook, so on_turn_end (the example's only entry point) has no way to release the env on a trainer-forced stop. Fixing it properly means adding a trajectory-end hook to the base scheduler and having GYMScheduler/OpenEnvScheduler close any surviving env — a change with broader blast radius across all schedulers, so I've left it out of this minimal fix and defer to maintainers on the preferred shape. Happy to follow up with that if you'd like.

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request fixes a KeyError and data loss issue in sudoku_scheduler.py by snapshotting rollout_infos before calling _close_and_remove(uuid), which cleans up the reward tracking dictionaries. The reviewer suggests also including other accumulated reward metrics (such as empty cell, valid move, repetition, progress, and correct rewards) in rollout_infos when a parse failure occurs on later turns, to prevent losing these metrics and to maintain consistency for downstream logging.

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.

[gemini-code-assist]

When a trajectory fails to parse on a later turn (i.e., current_turn > 1), it already has accumulated metrics in self._empty_cell_scores, self._valid_move_scores, etc. By only returning total_reward, step_rewards, and gym_done in rollout_infos, these accumulated metrics are lost. Additionally, downstream logging or metric aggregation might expect a consistent set of keys in rollout_infos and could fail or log incomplete data if these keys are missing.

We should include all the reward component metrics in rollout_infos here as well, using safe defaults/averages.

Suggested change

	rollout_infos = {
	'total_reward': self._total_rewards[uuid],
	'step_rewards': list(self._step_rewards.get(uuid, [])),
	'gym_done': True,
	}
	remaining = 81 - self._initial_filled.get(uuid, 0)
	progress_score = (self._max_filled.get(uuid, 0) - self._initial_filled.get(uuid, 0)) / remaining if remaining > 0 else 1.0
	rollout_infos = {
	'total_reward': self._total_rewards[uuid],
	'step_rewards': list(self._step_rewards.get(uuid, [])),
	'gym_done': True,
	'empty_cell_reward': sum(self._empty_cell_scores.get(uuid, [])) / max(len(self._empty_cell_scores.get(uuid, [])), 1),
	'valid_move_reward': sum(self._valid_move_scores.get(uuid, [])) / max(len(self._valid_move_scores.get(uuid, [])), 1),
	'repetition_reward': sum(self._repetition_scores.get(uuid, [])) / max(len(self._repetition_scores.get(uuid, [])), 1),
	'progress_reward': progress_score,
	'correct_reward': sum(self._correct_scores.get(uuid, [])) / max(len(self._correct_scores.get(uuid, [])), 1),
	}

[nantenT]

Thanks for the fix! The current change correctly addresses the rollout_infos loss issue by building the information before calling _close_and_remove, which avoids the state cleanup ordering problem.

I also want to share a possible workaround for the second issue mentioned in #9745 (OpenEnv session leak when the trajectory ends due to max_turns / length instead of env.done=True).

A minimal scheduler-level workaround could be adding the external termination conditions into on_turn_end. For example, in the Sudoku scheduler:

max_turns = 20

if done or current_turn >= max_turns - 1:
    await self._close_and_remove(uuid)
    done = True

return {
    "done": done,
    "rollout_infos": rollout_infos,
}

The same idea could be extended to finish_reason == "length" by tracking trajectory length/token usage inside the scheduler and treating it as another terminal condition.

This would provide a small, example-level workaround for users running multi-turn GRPO with OpenEnv, without requiring framework-level changes.

That said, I agree that the cleaner long-term solution would be introducing a scheduler lifecycle hook such as on_trajectory_end(uuid, reason), because trajectory termination can happen from multiple sources:

• environment returns done=True
• max_turns reached
• generation stops due to length
• exception / trainer interruption

A unified lifecycle callback would make cleanup behavior consistent across different schedulers.

Thanks again for the fix!

[hjh0119]

thanks