diff --git a/skyrl/train/generators/base.py b/skyrl/train/generators/base.py index 26d95868b9..d017b833df 100644 --- a/skyrl/train/generators/base.py +++ b/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]] diff --git a/skyrl/train/generators/skyrl_gym_generator.py b/skyrl/train/generators/skyrl_gym_generator.py index cc9bc95e80..56fd25915c 100644 --- a/skyrl/train/generators/skyrl_gym_generator.py +++ b/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, diff --git a/skyrl/train/generators/utils.py b/skyrl/train/generators/utils.py index 4c5228f189..98ae5c2138 100644 --- a/skyrl/train/generators/utils.py +++ b/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,6 +389,21 @@ 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]]], @@ -384,6 +411,9 @@ def get_rollout_metrics( 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) diff --git a/tests/train/generators/test_generator_output_utils.py b/tests/train/generators/test_generator_output_utils.py index 7546858bb8..827dd4510d 100644 --- a/tests/train/generators/test_generator_output_utils.py +++ b/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", diff --git a/tests/train/generators/test_skyrl_gym_generator.py b/tests/train/generators/test_skyrl_gym_generator.py index c29afbe920..b20d986797 100644 --- a/tests/train/generators/test_skyrl_gym_generator.py +++ b/tests/train/generators/test_skyrl_gym_generator.py @@ -1740,14 +1740,162 @@ def step(self, action): # Aggregate stats are present and computed over the per-prompt times. rollout_metrics = generator_output["rollout_metrics"] for key in ( - "generate/trajectory_completion_time_mean", - "generate/trajectory_completion_time_p90", - "generate/trajectory_completion_time_max", + "generate/trajectory_time_completion_mean", + "generate/trajectory_time_completion_p90", + "generate/trajectory_time_completion_max", ): assert key in rollout_metrics, f"missing metric {key}" expected = np.array(metrics_times, dtype=np.float64) - assert rollout_metrics["generate/trajectory_completion_time_mean"] == pytest.approx(np.mean(expected).item()) - assert rollout_metrics["generate/trajectory_completion_time_p90"] == pytest.approx( + assert rollout_metrics["generate/trajectory_time_completion_mean"] == pytest.approx(np.mean(expected).item()) + assert rollout_metrics["generate/trajectory_time_completion_p90"] == pytest.approx( np.percentile(expected, 90).item() ) - assert rollout_metrics["generate/trajectory_completion_time_max"] == pytest.approx(np.max(expected).item()) + assert rollout_metrics["generate/trajectory_time_completion_max"] == pytest.approx(np.max(expected).item()) + + # The llm/env/setup split is stored per step and replicated like the completion times above. + for key in ("trajectory_llm_times", "trajectory_env_times", "trajectory_env_setup_times"): + split_times = generator_output[key] + assert split_times is not None + assert len(split_times) == num_steps + assert split_times[0] == split_times[1] and split_times[2] == split_times[3] + + +@pytest.mark.asyncio +@patch("skyrl_gym.make") +async def test_llm_vs_env_time_split_metrics(mock_make, mock_tokenizer, mock_llm, mock_env_cfg): + """Only the total rollout time was previously recorded, which cannot distinguish an engine-bound + rollout from an environment-bound one. Uses an env deliberately slower than the engine, so a + swapped or one-sided attribution fails.""" + import asyncio + import time + + import numpy as np + + from skyrl.train.generators import utils as generator_utils + from skyrl.train.generators.base import TrajectoryID + + llm_sleep_s = 0.02 + env_sleep_s = 0.06 + setup_sleep_s = 0.03 + num_turns = 2 + + mock_tokenizer.eos_token_id = 4 + + def apply_chat_template_side_effect(messages, **kwargs): + if kwargs.get("tokenize", True): + return [201, 202] + return "".join([m.get("content", "") for m in messages]) + + mock_tokenizer.apply_chat_template.side_effect = apply_chat_template_side_effect + + async def llm_generate_side_effect(input_batch, model=None): + await asyncio.sleep(llm_sleep_s) + num = len(input_batch["prompt_token_ids"]) if "prompt_token_ids" in input_batch else len(input_batch["prompts"]) + return { + "responses": ["step"] * num, + "stop_reasons": ["stop"] * num, + "response_logprobs": None, + "response_ids": [[10, 11, 12, mock_tokenizer.eos_token_id] for _ in range(num)], + } + + mock_llm.generate = AsyncMock(side_effect=llm_generate_side_effect) + + class SlowEnv(BaseTextEnv): + def __init__(self): + super().__init__() + self.turns = 0 + + def init(self, prompt): + time.sleep(setup_sleep_s) + return prompt, {} + + def step(self, action): + time.sleep(env_sleep_s) + self.turns += 1 + if self.turns < num_turns: + return BaseTextEnvStepOutput( + observations=[{"role": "user", "content": "obs"}], reward=0.0, done=False, metadata={} + ) + return BaseTextEnvStepOutput(observations=[], reward=1.0, done=True, metadata={}) + + mock_make.side_effect = lambda *args, **kwargs: SlowEnv() + + # Run env.step on the executor as in production; a blocking step inline would stall the event + # loop and inflate a sibling trajectory's measured LLM time. + mock_env_cfg.max_env_workers = 4 + + cfg = GeneratorConfig() + cfg.sampling_params.max_generate_length = 50 + cfg.sampling_params.logprobs = None + cfg.apply_overlong_filtering = False + cfg.max_input_length = 512 + cfg.batched = False + cfg.max_turns = 10 + cfg.zero_reward_on_non_stop = False + cfg.use_conversation_multi_turn = True + cfg.chat_template = ChatTemplateConfig(source="name", name_or_path=None) + + generator = SkyRLGymGenerator( + generator_cfg=cfg, + skyrl_gym_cfg=mock_env_cfg, + inference_engine_client=mock_llm, + tokenizer=mock_tokenizer, + ) + generator.base_conversation_token_ids = [] + + num_trajectories = 2 + prompts = [[{"role": "user", "content": f"Q{i}?"}] for i in range(num_trajectories)] + input_batch: GeneratorInput = { + "prompts": prompts, + "env_extras": [{} for _ in prompts], + "env_classes": [mock_env_cfg.env_class for _ in prompts], + "trajectory_ids": [TrajectoryID(instance_id=f"uid{i}", repetition_id=0) for i in range(num_trajectories)], + } + + spy = MagicMock(side_effect=generator_utils.get_rollout_metrics) + with patch("skyrl.train.generators.skyrl_gym_generator.get_rollout_metrics", spy): + generator_output: GeneratorOutput = await generator.generate(input_batch) + + llm_times = spy.call_args.kwargs["trajectory_llm_times"] + env_times = spy.call_args.kwargs["trajectory_env_times"] + setup_times = spy.call_args.kwargs["trajectory_env_setup_times"] + assert llm_times is not None and env_times is not None and setup_times is not None + assert len(llm_times) == num_trajectories + assert len(env_times) == num_trajectories + assert len(setup_times) == num_trajectories + assert all(t >= setup_sleep_s for t in setup_times) + + # Each side is at least its per-turn sleep summed over turns, and neither swallows the other. + for llm_t, env_t in zip(llm_times, env_times): + assert llm_t >= llm_sleep_s * num_turns + assert env_t >= env_sleep_s * num_turns + # The environment is ~3x slower than the engine here, so a swapped attribution would flip this. + assert env_t > llm_t + + metrics = generator_output["rollout_metrics"] + assert metrics["generate/trajectory_time_llm_mean"] >= llm_sleep_s * num_turns + assert metrics["generate/trajectory_time_env_mean"] >= env_sleep_s * num_turns + assert metrics["generate/trajectory_time_env_setup_mean"] >= setup_sleep_s + assert metrics["generate/trajectory_time_other_mean"] >= 0.0 + + # env_sleep / (env_sleep + llm_sleep) = 0.75; allow generous headroom for scheduling overhead + # attributed to the engine wait, but it must clearly indicate an environment-bound rollout. + assert 0.5 < metrics["generate/frac_time_in_env"] < 1.0 + + # The bands never exceed the trajectory's end-to-end time; "other" is the exact remainder. + for llm_t, env_t, setup_t, e2e_t in zip( + llm_times, env_times, setup_times, generator_output["trajectory_generation_times"] + ): + assert llm_t + env_t + setup_t <= e2e_t + 1e-6 + + # Regression: concatenate_generator_outputs (every logging path) must recompute the split from + # the raw per-trajectory lists, not sum per-group p90/frac past the sample max and 1.0. + concatenated = generator_utils.concatenate_generator_outputs([generator_output, generator_output]) + concat_metrics = concatenated["rollout_metrics"] + assert 0.5 < concat_metrics["generate/frac_time_in_env"] < 1.0 + assert concat_metrics["generate/trajectory_time_llm_p90"] == pytest.approx(np.percentile(llm_times * 2, 90).item()) + assert concat_metrics["generate/trajectory_time_env_p90"] == pytest.approx(np.percentile(env_times * 2, 90).item()) + assert concat_metrics["generate/trajectory_time_env_setup_p90"] == pytest.approx( + np.percentile(setup_times * 2, 90).item() + ) + assert concat_metrics["generate/trajectory_time_other_mean"] >= 0.0