feat(generators): add generator.skip_failed_rollouts

[pranavraja99]

Closes #1613.

SkyRLGymGenerator.generate fans rollouts out through tqdm.gather, which inherits asyncio.gather's default semantics: the first exception from any rollout aborts the whole training step and throws away the trajectories that already finished. In a flaky multi-turn agentic setting (a network blip in a coding agent, say) one bad rollout wastes the entire step's compute.

generator.skip_failed_rollouts=true (default false) contains each rollout's exceptions instead. A failed rollout is replaced by a single-token placeholder trajectory with zero reward, a zeroed loss mask, and stop_reason="rollout_error". The placeholder preserves the batch size the trainer's data-parallel sharding expects, and since the stop reason is not "stop", the existing zero_reward_on_non_stop and apply_overlong_filtering knobs already treat it as a trajectory that should not contribute — no changes needed there.

Notes on the details that the issue's sketch left open:

• Reward shape. validate_generator_output requires every reward in a batch to be uniformly trajectory-level or token-level, and which one is in play is only known from the rollouts that succeeded. So a failing rollout returns a FailedRollout marker and substitution happens after the gather, once a successful output can be inspected. A fresh placeholder is built per failure so that later in-place edits to a reward or loss mask cannot alias across trajectories.
• All rollouts failing. The first exception is re-raised as the cause of a RuntimeError. A step consisting only of placeholders carries no training signal, and a systemic failure (engine down) should surface rather than be smoothed over.
• Exception, not BaseException. Cancellation and KeyboardInterrupt still propagate and tear down the batch.
• Unsupported combinations. Rejected at construction for batched=True (that path does not go through agent_loop, so the flag would silently do nothing) and for enable_return_routed_experts=True (a placeholder has no routed-expert indices to stand in for).

Environment cleanup

This also moves env.close() into agent_loop's existing finally. It previously ran only on the happy path, so a rollout raising mid-episode never released its environment. The leak was inconsequential while any failure killed the step, but it is not once failures are routine and the run continues — so it is in scope here rather than deferred. It is narrow (one env per failed agent_loop) and distinct from the trial/sandbox leakage in #1194. Happy to split it into its own PR if you'd prefer.

Testing

New CPU tests in tests/train/generators/test_skyrl_gym_generator.py:

• placeholder substitution keeps the successful rollout intact, and the result passes the trainer's real validate_generator_output
• placeholder reward follows token-level rewards, and carries one logprob per token when logprobs are requested
• all-rollouts-failed raises with the original exception as __cause__
• without the flag, a failed rollout still aborts the step (unchanged default)
• batched=True and enable_return_routed_experts=True are rejected
• agent_loop closes the env when the rollout raises (fails without the cleanup change)

uv run --isolated --extra skyrl-train --extra dev pytest tests/train/ tests/backends/skyrl_train/ --ignore=tests/backends/skyrl_train/gpu -m "not vllm" — 1463 passed, 7 skipped. Docs build (npm run build) passes.

Known behavior

e2e_time and time_splits are None on a placeholder, so a batch containing any failed rollout omits the trajectory completion-time metrics for that step. This follows the existing convention in generate of dropping those fields entirely rather than emitting a partially-populated list, and it seemed better than attributing a fabricated 0.0 to a rollout that never ran. Happy to change it if you'd rather see partial timings.

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request introduces the skip_failed_rollouts feature to SkyRLGymGenerator, allowing training steps to tolerate individual rollout failures by substituting them with single-token placeholder trajectories instead of aborting the entire batch. The review feedback highlights two main improvement areas: first, handling potential IndexError crashes when step_outputs is empty by safely resolving the first trajectory; second, tracking and preserving elapsed time (e2e_time) for failed rollouts so that batch-level completion-time metrics remain accurate.

[gemini-code-assist]

If a rollout succeeds but has an empty step_outputs (for example, if the initial prompt length exceeds max_input_length on the very first iteration, causing an immediate break in agent_loop), output.step_outputs will be empty. Accessing succeeded[0].step_outputs[0] directly will raise an IndexError. We should safely find the first non-empty trajectory to determine the reward shape.

        first_traj = None
        for output in succeeded:
            first_traj = self._first_trajectory(output)
            if first_traj is not None:
                break
        token_level_rewards = isinstance(first_traj.reward, list) if first_traj is not None else False

[gemini-code-assist]

Modify _first_trajectory to safely return None if step_outputs is empty, preventing potential IndexError crashes.

Suggested change

	def _first_trajectory(output: Union[TrajectoryOutput, StepWiseOutput]) -> TrajectoryOutput:
	"""The first per-step trajectory of a step-wise output, or the trajectory itself."""
	if isinstance(output, StepWiseOutput):
	return output.step_outputs[0]
	return output
	@staticmethod
	def _first_trajectory(output: Union[TrajectoryOutput, StepWiseOutput]) -> Optional[TrajectoryOutput]:
	"""The first per-step trajectory of a step-wise output, or the trajectory itself."""
	if isinstance(output, StepWiseOutput):
	return output.step_outputs[0] if output.step_outputs else None

[gemini-code-assist]

Add e2e_time to FailedRollout so we can track the elapsed time of failed rollouts and preserve the completion-time metrics for the successful rollouts in the batch.

Suggested change

	class FailedRollout:
	"""Marker returned in place of a rollout that raised under ``skip_failed_rollouts``."""
	exception: Exception
	class FailedRollout:
	"""Marker returned in place of a rollout that raised under ``skip_failed_rollouts``."""
	exception: Exception
	e2e_time: float = 0.0

[gemini-code-assist]

Measure the elapsed time of the failed rollout in _rollout_or_marker and store it in FailedRollout.

    async def _rollout_or_marker(
        self, index: int, rollout: Awaitable[Union[TrajectoryOutput, StepWiseOutput]]
    ) -> Union[TrajectoryOutput, StepWiseOutput, FailedRollout]:
        """Await a single rollout, turning a raised exception into a ``FailedRollout`` marker.

        Catches ``Exception`` rather than ``BaseException`` so that cancellation and keyboard
        interrupts still propagate and tear down the batch.
        """
        start_time = time.monotonic()
        try:
            return await rollout
        except Exception as exc:
            elapsed = time.monotonic() - start_time
            logger.warning(
                f"Rollout {index} failed with {type(exc).__name__}: {exc}. Substituting a placeholder "
                f"trajectory so the training step can complete."
            )
            return FailedRollout(exception=exc, e2e_time=elapsed)

[gemini-code-assist]

Pass e2e_time to _placeholder_rollout and assign it to the placeholder trajectory so that the batch's completion-time metrics are not completely discarded when a rollout fails.

    def _placeholder_rollout(
        self, token_level_rewards: bool, get_logprobs: bool, e2e_time: float = 0.0
    ) -> Union[TrajectoryOutput, StepWiseOutput]:
        """Build a single-token trajectory that contributes neither reward nor loss."""
        # `validate_generator_output` requires a non-empty response, and equal lengths for the
        # response, loss mask, token-level rewards and rollout logprobs.
        pad_token_id = self.tokenizer.pad_token_id
        if pad_token_id is None:
            pad_token_id = self.tokenizer.eos_token_id
        trajectory = TrajectoryOutput(
            response_ids=[pad_token_id],
            reward=[0.0] if token_level_rewards else 0.0,
            stop_reason=ROLLOUT_ERROR_STOP_REASON,
            loss_mask=[0],
            prompt_ids=[pad_token_id],
            rollout_logprobs=[0.0] if get_logprobs else None,
            env_metrics={},
            e2e_time=e2e_time,
        )
        if self.generator_cfg.step_wise_trajectories:
            return StepWiseOutput(step_outputs=[trajectory], e2e_time=e2e_time)
        return trajectory

[gemini-code-assist]

Pass the failed rollout's e2e_time to _placeholder_rollout during substitution.

Suggested change

	return [
	(
	self._placeholder_rollout(token_level_rewards, get_logprobs)
	if isinstance(output, FailedRollout)
	else output
	)
	for output in all_outputs
	]
	return [
	(
	self._placeholder_rollout(token_level_rewards, get_logprobs, output.e2e_time)
	if isinstance(output, FailedRollout)
	else output
	)
	for output in all_outputs
	]