[logging] Account for env setup and overhead in trajectory time

skyrl/train/generators/base.py

@@ -46,6 +46,11 @@ 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]]
+    # Split of ``trajectory_generation_times`` into engine-wait vs ``env.step()`` time (seconds),
+    # plus env construction/init time. None if any trajectory did not record it.
+    trajectory_llm_times: Optional[List[float]]
+    trajectory_env_times: Optional[List[float]]
+    trajectory_env_setup_times: Optional[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,12 @@ 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 time (seconds) awaiting the inference engine, summed over turns. None if untracked.
+    llm_time: Optional[float] = None
+    # Wall-clock time (seconds) in ``env.step()``, summed over turns. None if untracked.
+    env_time: Optional[float] = None
+    # Wall-clock time (seconds) constructing and initializing the env. None if untracked.
+    env_setup_time: Optional[float] = None
 
 
 @dataclass
@@ -66,6 +72,10 @@ 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
+    # Same split as TrajectoryOutput.llm_time / env_time / env_setup_time.
+    llm_time: Optional[float] = None
+    env_time: Optional[float] = None
+    env_setup_time: Optional[float] = None
 
 
 @dataclass
@@ -140,6 +150,12 @@ def get_turn_rollout_logprobs(self) -> Optional[List[float]]:
         return self.output_logprobs + [0.0] * len(self.obs_ids)
 
 
+def _optional_times(outputs, attr: str) -> Optional[List[float]]:
+    """One value per trajectory, or None if any trajectory did not record ``attr``."""
+    values = [getattr(o, attr, None) for o in outputs]
+    return None if any(v is None for v in values) else values
+
+
 class SkyRLGymGenerator(GeneratorInterface):
     def __init__(
         self,
@@ -321,6 +337,10 @@ async def agent_loop(
             rollout_logprobs: Optional[List[float]]
         """
         agent_loop_start_time = time.monotonic()
+        # Engine wait vs. env.step() time, accumulated across turns. The gap to e2e_time is
+        # tokenization, chat templating, and event-loop scheduling.
+        llm_time_s = 0.0
+        env_time_s = 0.0
 
         session_id = (
             f"{trajectory_id.instance_id}_{trajectory_id.repetition_id}" if trajectory_id is not None else uuid4().hex
@@ -336,6 +356,7 @@ async def agent_loop(
             env_extras = self._setup_env_extras(env_class, env_extras, sampling_params, trajectory_id)
 
             env_config = getattr(self.skyrl_gym_cfg, env_class, dict())
+            env_setup_start_time = time.monotonic()
             env = skyrl_gym.make(env_class, env_config=env_config, extras=env_extras)
 
             # Instantiate chat_history and chat_end_index, which are only used if `retokenize_chat_history==True`.
@@ -344,6 +365,7 @@ async def agent_loop(
 
             # init() returns the first prompt to be given to the model, and optional metadata dict
             chat_history, _ = await self._run_in_executor_if_available(env.init, chat_history)
+            env_setup_time_s = time.monotonic() - env_setup_start_time
             initial_chat_history_length = len(chat_history)
             initial_input_ids = self.tokenizer.apply_chat_template(
                 chat_history,
@@ -409,7 +431,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)
+                llm_time_s += 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 +467,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)
+                env_time_s += 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 +632,9 @@ async def agent_loop(
                 trajectory_id,
             )
             agent_loop_output.e2e_time = time.monotonic() - agent_loop_start_time
+            agent_loop_output.llm_time = llm_time_s
+            agent_loop_output.env_time = env_time_s
+            agent_loop_output.env_setup_time = env_setup_time_s
             return agent_loop_output
 
         finally:
@@ -868,11 +897,10 @@ async def generate(self, input_batch: GeneratorInput, disable_tqdm: bool = False
             disable=disable_tqdm,
         )
         # Per-trajectory end-to-end generation times (one entry per prompt, preserving input order).
-        # ``e2e_time`` is optional for agent loops; if any trajectory did not record it, we omit the
-        # field entirely rather than emit a partially-populated list.
-        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_generation_times_per_prompt = _optional_times(all_outputs, "e2e_time")
+        trajectory_llm_times_per_prompt = _optional_times(all_outputs, "llm_time")
+        trajectory_env_times_per_prompt = _optional_times(all_outputs, "env_time")
+        trajectory_env_setup_times_per_prompt = _optional_times(all_outputs, "env_setup_time")
 
         if self.generator_cfg.step_wise_trajectories:
             responses = []
@@ -885,6 +913,9 @@ async def generate(self, input_batch: GeneratorInput, disable_tqdm: bool = False
             out_trajectory_ids = []
             out_env_classes = []
             out_trajectory_generation_times = []
+            out_trajectory_llm_times = []
+            out_trajectory_env_times = []
+            out_trajectory_env_setup_times = []
             for i, output in enumerate(all_outputs):
                 for j, step_output in enumerate(output.step_outputs):
                     responses.append(step_output.response_ids)
@@ -896,11 +927,20 @@ 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_trajectory_llm_times.append(getattr(output, "llm_time", None))
+                    out_trajectory_env_times.append(getattr(output, "env_time", None))
+                    out_trajectory_env_setup_times.append(getattr(output, "env_setup_time", None))
             # Keep aligned with the per-prompt None handling:
             if not trajectory_generation_times_per_prompt:
                 out_trajectory_generation_times = None
+            if not trajectory_llm_times_per_prompt:
+                out_trajectory_llm_times = None
+            if not trajectory_env_times_per_prompt:
+                out_trajectory_env_times = None
+            if not trajectory_env_setup_times_per_prompt:
+                out_trajectory_env_setup_times = None
             env_classes = out_env_classes
         else:
             responses = [output.response_ids for output in all_outputs]
@@ -913,6 +953,9 @@ 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_llm_times = trajectory_llm_times_per_prompt
+            out_trajectory_env_times = trajectory_env_times_per_prompt
+            out_trajectory_env_setup_times = trajectory_env_setup_times_per_prompt
 
         has_vision_features = any(getattr(output, "pixel_values", None) is not None for output in all_outputs)
         pixel_values = (
@@ -953,6 +996,9 @@ 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_llm_times=trajectory_llm_times_per_prompt,
+            trajectory_env_times=trajectory_env_times_per_prompt,
+            trajectory_env_setup_times=trajectory_env_setup_times_per_prompt,
         )
 
         if self.generator_cfg.zero_reward_on_non_stop:
@@ -974,6 +1020,9 @@ 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_llm_times": out_trajectory_llm_times,
+            "trajectory_env_times": out_trajectory_env_times,
+            "trajectory_env_setup_times": out_trajectory_env_setup_times,
             "rollout_expert_indices": rollout_expert_indices,
             "is_last_step": is_last_step,
             "env_metrics": env_metrics,

skyrl/train/generators/utils.py

@@ -240,6 +240,20 @@ def _flatten_field(generator_outputs: List[GeneratorOutput], key: str) -> list:
     return flat
 
 
+def _concat_field(generator_outputs: List[GeneratorOutput], key: str) -> Optional[list]:
+    """Flatten an optional per-trajectory field, keyed off the first output (None if absent)."""
+    if generator_outputs[0].get(key) is None:
+        return None
+    return _flatten_field(generator_outputs, key)
+
+
+def _last_step_only(values: Optional[list], is_last_step: Optional[List[bool]]) -> Optional[list]:
+    """Keep only last-step entries when step-wise, so one trajectory contributes one value."""
+    if values is None or not is_last_step:
+        return values
+    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 +278,12 @@ 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_field(generator_outputs, "stop_reasons"),
+        "rollout_logprobs": _concat_field(generator_outputs, "rollout_logprobs"),
+        "trajectory_generation_times": _concat_field(generator_outputs, "trajectory_generation_times"),
+        "trajectory_llm_times": _concat_field(generator_outputs, "trajectory_llm_times"),
+        "trajectory_env_times": _concat_field(generator_outputs, "trajectory_env_times"),
+        "trajectory_env_setup_times": _concat_field(generator_outputs, "trajectory_env_setup_times"),
     }
 
     # propagate additional keys with list values as-is
@@ -284,19 +293,22 @@ 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_llm_times = _last_step_only(result.get("trajectory_llm_times"), is_last_step)
+    trajectory_env_times = _last_step_only(result.get("trajectory_env_times"), is_last_step)
+    trajectory_env_setup_times = _last_step_only(result.get("trajectory_env_setup_times"), 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_llm_times=trajectory_llm_times,
+        trajectory_env_times=trajectory_env_times,
+        trajectory_env_setup_times=trajectory_env_setup_times,
     )
 
     # Preserve generator-specific metrics from per-group rollout_metrics. get_rollout_metrics only
@@ -377,13 +389,31 @@ 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]]) -> Optional[float]:
+    """Add mean/p90/max for a per-trajectory time list; returns its sum, or None if absent."""
+    if not times:
+        return None
+    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(),
+        }
+    )
+    return np.sum(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_llm_times: Optional[List[float]] = None,
+    trajectory_env_times: Optional[List[float]] = None,
+    trajectory_env_setup_times: Optional[List[float]] = None,
 ):
     """
     Computes rollout metrics including token statistics and optional environment-specific metrics.
@@ -396,6 +426,12 @@ 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_llm_times: Optional per-trajectory time (seconds) spent awaiting the inference
+            engine, summed over turns
+        trajectory_env_times: Optional per-trajectory time (seconds) spent in ``env.step()``, summed
+            over turns
+        trajectory_env_setup_times: Optional per-trajectory time (seconds) spent constructing and
+            initializing the env
 
     Returns:
         Dictionary of aggregated metrics
@@ -450,15 +486,26 @@ 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)
+    llm_sum = _add_time_stats(rollout_metrics, "llm", trajectory_llm_times)
+    env_sum = _add_time_stats(rollout_metrics, "env", trajectory_env_times)
+    _add_time_stats(rollout_metrics, "env_setup", trajectory_env_setup_times)
+
+    if llm_sum is not None and env_sum is not None and llm_sum + env_sum > 0:
+        # Time-weighted (sum over sum), not a mean of per-trajectory ratios, so long trajectories
+        # count proportionally.
+        rollout_metrics["generate/frac_time_in_env"] = env_sum / (llm_sum + env_sum)
+
+    # Remainder of e2e not attributed to the engine, env steps, or env setup. Tokenization, chat
+    # templating, output assembly, env teardown, and event-loop scheduling.
+    if trajectory_completion_times and trajectory_llm_times and trajectory_env_times and trajectory_env_setup_times:
+        other = (
+            np.array(trajectory_completion_times, dtype=np.float64)
+            - np.array(trajectory_llm_times, dtype=np.float64)
+            - np.array(trajectory_env_times, dtype=np.float64)
+            - np.array(trajectory_env_setup_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

@@ -32,6 +32,9 @@ def test_generator_output_concatenation():
         # optional but present in the signature
         "trajectory_ids",
         "trajectory_generation_times",
+        "trajectory_llm_times",
+        "trajectory_env_times",
+        "trajectory_env_setup_times",
         "is_last_step",
         "env_metrics",
         "pixel_values",