[logging] Split rollout trajectory time into inference-engine wait vs env.step

skyrl/train/generators/base.py

@@ -46,6 +46,10 @@ class GeneratorOutput(TypedDict):
     # trajectory in the input batch (i.e. per ``agent_loop`` call). Used by the fully
     # async trainer to compute per-group / intra-group completion-time metrics.
     trajectory_generation_times: Optional[List[float]]
+    # Engine and env time splits of ``trajectory_generation_times``, one list entry per trajectory,
+    # e.g. {"llm": [...], "env": [...]}. trajectory_time_splits is None if any trajectory did not
+    # record its split.
+    trajectory_time_splits: Optional[Dict[str, List[float]]]
     rollout_expert_indices: Optional[List[List[List[List[int]]]]]  # [batch_size, seq_len, layer_num, topk]
     # Applicable only for step-wise training
     is_last_step: Optional[List[bool]]

skyrl/train/generators/skyrl_gym_generator.py

@@ -56,6 +56,9 @@ class TrajectoryOutput:
     # End-to-end wall-clock time (seconds) to generate this trajectory. Optional: agent loops may
     # leave this as None if they do not track timing.
     e2e_time: Optional[float] = None
+    # Wall-clock seconds spent in each phase. "llm" is time in inference-engine calls, "env" is
+    # time in env.step. Field is None if any loop did not record a split.
+    time_splits: Optional[Dict[str, float]] = None
 
 
 @dataclass
@@ -66,6 +69,9 @@ class StepWiseOutput:
     # End-to-end wall-clock time (seconds) to generate this trajectory. Optional: agent loops may
     # leave this as None if they do not track timing.
     e2e_time: Optional[float] = None
+    # Wall-clock seconds spent in each phase. "llm" is time in inference-engine calls, "env" is
+    # time in env.step. Field is None if any loop did not record a split.
+    time_splits: Optional[Dict[str, float]] = None
 
 
 @dataclass
@@ -140,6 +146,13 @@ def get_turn_rollout_logprobs(self) -> Optional[List[float]]:
         return self.output_logprobs + [0.0] * len(self.obs_ids)
 
 
+def _split_lists(time_splits: List[Optional[Dict[str, float]]]) -> Optional[Dict[str, List[float]]]:
+    """Per-component lists from per-trajectory time splits, or None if any trajectory lacks them."""
+    if not time_splits or any(s is None for s in time_splits):
+        return None
+    return {name: [s[name] for s in time_splits] for name in time_splits[0]}
+
+
 class SkyRLGymGenerator(GeneratorInterface):
     def __init__(
         self,
@@ -321,6 +334,7 @@ async def agent_loop(
             rollout_logprobs: Optional[List[float]]
         """
         agent_loop_start_time = time.monotonic()
+        time_splits = {"llm": 0.0, "env": 0.0}
 
         session_id = (
             f"{trajectory_id.instance_id}_{trajectory_id.repetition_id}" if trajectory_id is not None else uuid4().hex
@@ -409,7 +423,9 @@ async def agent_loop(
                     sampling_params=sampling_params,
                     cache_salt=cache_salt,
                 )
+                llm_call_start_time = time.monotonic()
                 engine_output = await self.inference_engine_client.generate(engine_input, model=self.policy_model_name)
+                time_splits["llm"] += time.monotonic() - llm_call_start_time
                 output = engine_output["responses"][0]
                 output_ids = engine_output["response_ids"][0]
                 stop_reason = engine_output["stop_reasons"][0]
@@ -443,7 +459,9 @@ async def agent_loop(
                         added_eos = True
 
                 # 2. Environment step
+                env_step_start_time = time.monotonic()
                 env_step_output: BaseTextEnvStepOutput = await self._run_in_executor_if_available(env.step, output)
+                time_splits["env"] += time.monotonic() - env_step_start_time
                 new_obs = env_step_output["observations"]
                 step_reward: float = env_step_output["reward"]
                 agent_loop_state.done = env_step_output["done"]
@@ -606,6 +624,7 @@ async def agent_loop(
                 trajectory_id,
             )
             agent_loop_output.e2e_time = time.monotonic() - agent_loop_start_time
+            agent_loop_output.time_splits = time_splits
             return agent_loop_output
 
         finally:
@@ -873,6 +892,7 @@ async def generate(self, input_batch: GeneratorInput, disable_tqdm: bool = False
         trajectory_generation_times_per_prompt = [getattr(output, "e2e_time", None) for output in all_outputs]
         if any(t is None for t in trajectory_generation_times_per_prompt):
             trajectory_generation_times_per_prompt = None
+        trajectory_time_splits_per_prompt = _split_lists([getattr(o, "time_splits", None) for o in all_outputs])
 
         if self.generator_cfg.step_wise_trajectories:
             responses = []
@@ -885,6 +905,7 @@ async def generate(self, input_batch: GeneratorInput, disable_tqdm: bool = False
             out_trajectory_ids = []
             out_env_classes = []
             out_trajectory_generation_times = []
+            out_step_time_splits = []
             for i, output in enumerate(all_outputs):
                 for j, step_output in enumerate(output.step_outputs):
                     responses.append(step_output.response_ids)
@@ -896,11 +917,13 @@ async def generate(self, input_batch: GeneratorInput, disable_tqdm: bool = False
                     is_last_step.append(j == len(output.step_outputs) - 1)
                     out_trajectory_ids.append(trajectory_ids[i])
                     out_env_classes.append(env_classes[i])
-                    # For trajectory completion per turn we just use the trajectory level e2e time
+                    # For trajectory completion per turn we just use the trajectory level times
                     out_trajectory_generation_times.append(getattr(output, "e2e_time", None))
+                    out_step_time_splits.append(getattr(output, "time_splits", None))
             # Keep aligned with the per-prompt None handling:
             if not trajectory_generation_times_per_prompt:
                 out_trajectory_generation_times = None
+            out_trajectory_time_splits = _split_lists(out_step_time_splits)
             env_classes = out_env_classes
         else:
             responses = [output.response_ids for output in all_outputs]
@@ -913,6 +936,7 @@ async def generate(self, input_batch: GeneratorInput, disable_tqdm: bool = False
             out_trajectory_ids = None
             # One time per trajectory, already aligned 1:1 with responses (None if not all recorded).
             out_trajectory_generation_times = trajectory_generation_times_per_prompt
+            out_trajectory_time_splits = trajectory_time_splits_per_prompt
 
         has_vision_features = any(getattr(output, "pixel_values", None) is not None for output in all_outputs)
         pixel_values = (
@@ -953,6 +977,7 @@ async def generate(self, input_batch: GeneratorInput, disable_tqdm: bool = False
             # NOTE: we only use trajectory completion times per prompt for
             # metrics, to avoid duplicate entries with step-wise training
             trajectory_completion_times=trajectory_generation_times_per_prompt,
+            trajectory_time_splits=trajectory_time_splits_per_prompt,
         )
 
         if self.generator_cfg.zero_reward_on_non_stop:
@@ -974,6 +999,7 @@ async def generate(self, input_batch: GeneratorInput, disable_tqdm: bool = False
             "trajectory_ids": out_trajectory_ids,
             # NOTE: for completion metrics, we output the completion time
             "trajectory_generation_times": out_trajectory_generation_times,
+            "trajectory_time_splits": out_trajectory_time_splits,
             "rollout_expert_indices": rollout_expert_indices,
             "is_last_step": is_last_step,
             "env_metrics": env_metrics,

skyrl/train/generators/utils.py

@@ -240,6 +240,31 @@ def _flatten_field(generator_outputs: List[GeneratorOutput], key: str) -> list:
     return flat
 
 
+def _concat_optional_field(
+    generator_outputs: List[GeneratorOutput], key: str
+) -> Optional[Union[list, Dict[str, list]]]:
+    """Concatenate an optional per-trajectory field across outputs. None if any output lacks it.
+    Handles a flat list or a dict-of-lists, concatenating each component."""
+    values = [go.get(key) for go in generator_outputs]
+    if any(v is None for v in values):
+        return None
+    if isinstance(values[0], dict):
+        return {name: [t for v in values for t in v[name]] for name in values[0]}
+    return _flatten_field(generator_outputs, key)
+
+
+def _last_step_only(
+    values: Optional[Union[list, Dict[str, list]]], is_last_step: Optional[List[bool]]
+) -> Optional[Union[list, Dict[str, list]]]:
+    """Keep only last-step entries when step-wise, so one trajectory contributes one value.
+    Handles a flat list or a dict-of-lists. No-op when not step-wise."""
+    if values is None or not is_last_step:
+        return values
+    if isinstance(values, dict):
+        return {name: _last_step_only(v, is_last_step) for name, v in values.items()}
+    return [v for v, last in zip(values, is_last_step) if last]
+
+
 def concatenate_generator_outputs(generator_outputs: List[GeneratorOutput], step_wise: bool = False) -> GeneratorOutput:
     """
     Concatenate the generator outputs of multiple batches. Then validate the concatenated result.
@@ -264,17 +289,10 @@ def concatenate_generator_outputs(generator_outputs: List[GeneratorOutput], step
         "response_ids": _flatten_field(generator_outputs, "response_ids"),
         "rewards": _flatten_field(generator_outputs, "rewards"),
         "loss_masks": _flatten_field(generator_outputs, "loss_masks"),
-        "stop_reasons": (
-            _flatten_field(generator_outputs, "stop_reasons") if first.get("stop_reasons") is not None else None
-        ),
-        "rollout_logprobs": (
-            _flatten_field(generator_outputs, "rollout_logprobs") if first.get("rollout_logprobs") is not None else None
-        ),
-        "trajectory_generation_times": (
-            _flatten_field(generator_outputs, "trajectory_generation_times")
-            if first.get("trajectory_generation_times") is not None
-            else None
-        ),
+        "stop_reasons": _concat_optional_field(generator_outputs, "stop_reasons"),
+        "rollout_logprobs": _concat_optional_field(generator_outputs, "rollout_logprobs"),
+        "trajectory_generation_times": _concat_optional_field(generator_outputs, "trajectory_generation_times"),
+        "trajectory_time_splits": _concat_optional_field(generator_outputs, "trajectory_time_splits"),
     }
 
     # propagate additional keys with list values as-is
@@ -284,19 +302,18 @@ def concatenate_generator_outputs(generator_outputs: List[GeneratorOutput], step
     for key in additional_keys:
         result[key] = _flatten_field(generator_outputs, key)
 
-    # With step-wise training, only use the trajectory generation time from the last step
-    trajectory_generation_times = result.get("trajectory_generation_times")
-    if step_wise and trajectory_generation_times and result.get("is_last_step"):
-        trajectory_generation_times = [
-            t for t, is_last_step in zip(trajectory_generation_times, result.get("is_last_step")) if is_last_step
-        ]
+    # With step-wise training each trajectory spans multiple rows; keep only its last-step timing.
+    is_last_step = result.get("is_last_step") if step_wise else None
+    trajectory_generation_times = _last_step_only(result.get("trajectory_generation_times"), is_last_step)
+    trajectory_time_splits = _last_step_only(result.get("trajectory_time_splits"), is_last_step)
 
-    # Re-aggregate rollout metrics
+    # Re-aggregate rollout metrics; the extra_keys fallback below cannot aggregate a p90 or a ratio.
     rollout_metrics = get_rollout_metrics(
         result["response_ids"],
         result["rewards"],
         loss_masks=result.get("loss_masks"),
         trajectory_completion_times=trajectory_generation_times,
+        trajectory_time_splits=trajectory_time_splits,
     )
 
     # Preserve generator-specific metrics from per-group rollout_metrics. get_rollout_metrics only
@@ -377,13 +394,28 @@ def compute_turn_token_counts(loss_masks: List[List[int]]) -> List[int]:
     return turn_token_counts
 
 
+def _add_time_stats(rollout_metrics: Dict[str, Any], name: str, times: Optional[List[float]]) -> None:
+    """Add mean/p90/max stats for a per-trajectory time list under generate/trajectory_time_<name>_*."""
+    if not times:
+        return
+    arr = np.array(times, dtype=np.float64)
+    rollout_metrics.update(
+        {
+            f"generate/trajectory_time_{name}_mean": np.mean(arr).item(),
+            f"generate/trajectory_time_{name}_p90": np.percentile(arr, 90).item(),
+            f"generate/trajectory_time_{name}_max": np.max(arr).item(),
+        }
+    )
+
+
 def get_rollout_metrics(
     responses: List[List[int]],
     rewards: Union[List[float], List[List[float]]],
     env_metrics: Optional[List[Dict[str, Any]]] = None,
     env_classes: Optional[List[str]] = None,
     loss_masks: Optional[List[List[int]]] = None,
     trajectory_completion_times: Optional[List[float]] = None,
+    trajectory_time_splits: Optional[Dict[str, List[float]]] = None,
 ):
     """
     Computes rollout metrics including token statistics and optional environment-specific metrics.
@@ -396,6 +428,8 @@ def get_rollout_metrics(
         loss_masks: Optional list of per-token loss masks; used to compute assistant-only token counts
         trajectory_completion_times: Optional per-trajectory end-to-end generation times (seconds);
             used to compute aggregate trajectory completion-time stats (mean / p90 / max)
+        trajectory_time_splits: Optional per-component split of the completion times, e.g.
+            {"llm": [...], "env": [...]}, one entry per trajectory
 
     Returns:
         Dictionary of aggregated metrics
@@ -450,15 +484,16 @@ def get_rollout_metrics(
                 }
             )
 
-    if trajectory_completion_times:
-        completion_times_arr = np.array(trajectory_completion_times, dtype=np.float64)
-        rollout_metrics.update(
-            {
-                "generate/trajectory_completion_time_mean": np.mean(completion_times_arr).item(),
-                "generate/trajectory_completion_time_p90": np.percentile(completion_times_arr, 90).item(),
-                "generate/trajectory_completion_time_max": np.max(completion_times_arr).item(),
-            }
-        )
+    _add_time_stats(rollout_metrics, "completion", trajectory_completion_times)
+    for name, times in (trajectory_time_splits or {}).items():
+        _add_time_stats(rollout_metrics, name, times)
+
+    if trajectory_completion_times and trajectory_time_splits:
+        # Completion time not attributed to a named split, e.g. tokenization and chat templating.
+        other = np.array(trajectory_completion_times, dtype=np.float64)
+        for times in trajectory_time_splits.values():
+            other = other - np.array(times, dtype=np.float64)
+        _add_time_stats(rollout_metrics, "other", other.tolist())
 
     if env_metrics is not None and env_classes is not None:
         env_to_metrics = defaultdict(list)

tests/train/generators/test_generator_output_utils.py

@@ -12,6 +12,7 @@
     compute_turn_token_counts,
     concatenate_generator_outputs,
     get_metrics_from_generator_output,
+    get_rollout_metrics,
     merge_stepwise_output,
 )
 from skyrl.train.utils.utils import validate_cfg
@@ -32,6 +33,7 @@ def test_generator_output_concatenation():
         # optional but present in the signature
         "trajectory_ids",
         "trajectory_generation_times",
+        "trajectory_time_splits",
         "is_last_step",
         "env_metrics",
         "pixel_values",
@@ -93,6 +95,73 @@ def test_generator_output_concatenation():
         np.testing.assert_allclose(concatenated_output["rollout_metrics"][key], value)
 
 
+def test_time_split_rollout_metrics():
+    metrics = get_rollout_metrics(
+        responses=[[1, 2]] * 4,
+        rewards=[1.0] * 4,
+        trajectory_completion_times=[10.0, 20.0, 30.0, 40.0],
+        trajectory_time_splits={"llm": [4.0, 8.0, 12.0, 16.0], "env": [5.0, 10.0, 15.0, 20.0]},
+    )
+    assert metrics["generate/trajectory_time_llm_mean"] == 10.0
+    assert metrics["generate/trajectory_time_llm_p90"] == pytest.approx(np.percentile([4.0, 8.0, 12.0, 16.0], 90))
+    assert metrics["generate/trajectory_time_llm_max"] == 16.0
+    assert metrics["generate/trajectory_time_env_mean"] == 12.5
+    # "other" is the exact per-trajectory remainder, here [1.0, 2.0, 3.0, 4.0].
+    assert metrics["generate/trajectory_time_other_mean"] == pytest.approx(2.5)
+    assert metrics["generate/trajectory_time_other_max"] == pytest.approx(4.0)
+
+
+def test_time_splits_concatenation():
+    def make_output(times, splits) -> GeneratorOutput:
+        return {
+            "prompt_token_ids": [[1]] * len(times),
+            "response_ids": [[1, 2]] * len(times),
+            "rewards": [1.0] * len(times),
+            "loss_masks": [[1, 1]] * len(times),
+            "stop_reasons": ["stop"] * len(times),
+            "rollout_logprobs": None,
+            "trajectory_generation_times": times,
+            "trajectory_time_splits": splits,
+        }
+
+    out1 = make_output([10.0, 20.0], {"llm": [4.0, 8.0], "env": [5.0, 10.0]})
+    out2 = make_output([30.0, 40.0], {"llm": [12.0, 16.0], "env": [15.0, 20.0]})
+    concatenated = concatenate_generator_outputs([out1, out2])
+
+    assert concatenated["trajectory_time_splits"] == {"llm": [4.0, 8.0, 12.0, 16.0], "env": [5.0, 10.0, 15.0, 20.0]}
+    # Aggregates are recomputed over the combined sample, not combined from per-group aggregates.
+    concat_metrics = concatenated["rollout_metrics"]
+    assert concat_metrics["generate/trajectory_time_llm_p90"] == pytest.approx(
+        np.percentile([4.0, 8.0, 12.0, 16.0], 90)
+    )
+    assert concat_metrics["generate/trajectory_time_other_mean"] == pytest.approx(2.5)
+
+
+def test_time_splits_concatenation_partial_is_none():
+    """A batch that recorded no splits drops the whole time-split aggregate to None."""
+
+    def make_output(times, splits) -> GeneratorOutput:
+        return {
+            "prompt_token_ids": [[1]] * len(times),
+            "response_ids": [[1, 2]] * len(times),
+            "rewards": [1.0] * len(times),
+            "loss_masks": [[1, 1]] * len(times),
+            "stop_reasons": ["stop"] * len(times),
+            "rollout_logprobs": None,
+            "trajectory_generation_times": times,
+            "trajectory_time_splits": splits,
+        }
+
+    recorded = make_output([10.0, 20.0], {"llm": [4.0, 8.0], "env": [5.0, 10.0]})
+    missing = make_output([30.0, 40.0], None)
+    # Both orders: the None batch must not depend on being first or last.
+    for outputs in ([recorded, missing], [missing, recorded]):
+        concatenated = concatenate_generator_outputs(outputs)
+        assert concatenated["trajectory_time_splits"] is None
+        # generation_times has its own None-ness, so it still concatenates fully.
+        assert sorted(concatenated["trajectory_generation_times"]) == [10.0, 20.0, 30.0, 40.0]
+
+
 def test_get_metrics_from_generator_output():
     # Per trajectory rewards, where rewards are List[float]
     generator_output: GeneratorOutput = {