Add opt-in --warmup_steps to exclude benchmark cold-start

Add a warmup_steps parameter to EnvironmentStepTimingRecorder that drops
the first N ``env.step()`` calls (CUDA graph capture, lazy kernel
compilation) from the recorded environment-step timing. In serialized
mode the skipped steps' simulation call-count and time are rolled back so
the recorded aggregate stays consistent.

Expose it as an opt-in ``--warmup_steps`` flag (default 0, validated
non-negative, recorded in metadata) on the eight RL train/play benchmark
adapters. ``runtime.py`` is unchanged: its ``--warmup_frames`` already
excludes the startup step before the recorder is installed.

Add recorder unit tests for the host-return and serialized exclusion
paths and a parse-level test rejecting a negative ``--warmup_steps``.

Author: AntoineRichard

scripts/benchmarks/rl_games/benchmark_rl_games_play.py

@@ -63,6 +63,12 @@ def _parse_args(argv: list[str]):
         action="store_true",
         help="Measure a serialized synchronized simulation and outside-simulation step breakdown.",
     )
+    parser.add_argument(
+        "--warmup_steps",
+        type=int,
+        default=0,
+        help="Exclude the first N env.step() calls from environment-step timing (cold-start). Opt-in; default 0.",
+    )
     parser.add_argument(
         "--benchmark_formatter",
         type=str,
@@ -78,6 +84,8 @@ def _parse_args(argv: list[str]):
     args_cli, remaining_args = setup_preset_cli(parser, argv)
     if args_cli.num_frames <= 0:
         parser.error("--num_frames must be greater than zero")
+    if args_cli.warmup_steps < 0:
+        parser.error("--warmup_steps must be non-negative")
     sys.argv = [sys.argv[0]] + remaining_args
 
     return args_cli, remaining_args
@@ -170,6 +178,7 @@ def run(argv: list[str]) -> None:
                             "serialized_synchronized" if args_cli.measure_synchronized_step_breakdown else "host_return"
                         ),
                     },
+                    {"name": "environment_step_warmup_steps", "data": args_cli.warmup_steps},
                     {"name": "presets", "data": ",".join(cfg.presets)},
                 ]
             },
@@ -231,7 +240,9 @@ def policy(obs):
             return agent.get_action(obs, is_deterministic=agent.is_deterministic)
 
         environment_step_timer = stepping.EnvironmentStepTimingRecorder(
-            env, measure_synchronized_step_breakdown=args_cli.measure_synchronized_step_breakdown
+            env,
+            measure_synchronized_step_breakdown=args_cli.measure_synchronized_step_breakdown,
+            warmup_steps=args_cli.warmup_steps,
         )
         with environment_step_timer, BenchmarkMonitor(benchmark, interval=1.0):
             step_times, reward, ep_length, success_rate = stepping.run_play_loop(env, policy, args_cli.num_frames)

scripts/benchmarks/rl_games/benchmark_rl_games_train.py

@@ -55,6 +55,12 @@ def _parse_args(argv: list[str]):
         action="store_true",
         help="Measure a serialized synchronized simulation and outside-simulation step breakdown.",
     )
+    parser.add_argument(
+        "--warmup_steps",
+        type=int,
+        default=0,
+        help="Exclude the first N env.step() calls from environment-step timing (cold-start). Opt-in; default 0.",
+    )
     parser.add_argument(
         "--benchmark_formatter",
         type=str,
@@ -88,6 +94,8 @@ def _parse_args(argv: list[str]):
     args_cli, remaining_args = setup_preset_cli(parser, argv)
     if args_cli.max_iterations is not None and args_cli.max_iterations <= 0:
         parser.error("--max_iterations must be greater than zero")
+    if args_cli.warmup_steps < 0:
+        parser.error("--warmup_steps must be non-negative")
     enable_cameras_for_video(args_cli)
     sys.argv = [sys.argv[0]] + remaining_args
 
@@ -189,6 +197,7 @@ def run(argv: list[str]) -> None:
                             "serialized_synchronized" if args_cli.measure_synchronized_step_breakdown else "host_return"
                         ),
                     },
+                    {"name": "environment_step_warmup_steps", "data": args_cli.warmup_steps},
                     {"name": "presets", "data": ",".join(cfg.presets)},
                 ]
             },
@@ -234,7 +243,9 @@ def run(argv: list[str]) -> None:
         runner.reset()
 
         environment_step_timer = stepping.EnvironmentStepTimingRecorder(
-            env, measure_synchronized_step_breakdown=args_cli.measure_synchronized_step_breakdown
+            env,
+            measure_synchronized_step_breakdown=args_cli.measure_synchronized_step_breakdown,
+            warmup_steps=args_cli.warmup_steps,
         )
         with environment_step_timer, BenchmarkMonitor(benchmark, interval=1.0):
             runner.run({"train": True, "play": False, "sigma": None})

scripts/benchmarks/rsl_rl/benchmark_rsl_rl_play.py

@@ -65,6 +65,12 @@ def _parse_args(argv: list[str]) -> tuple[argparse.Namespace, list[str]]:
         action="store_true",
         help="Measure a serialized synchronized simulation and outside-simulation step breakdown.",
     )
+    parser.add_argument(
+        "--warmup_steps",
+        type=int,
+        default=0,
+        help="Exclude the first N env.step() calls from environment-step timing (cold-start). Opt-in; default 0.",
+    )
     parser.add_argument(
         "--benchmark_formatter",
         type=str,
@@ -80,6 +86,8 @@ def _parse_args(argv: list[str]) -> tuple[argparse.Namespace, list[str]]:
     args, remaining = setup_preset_cli(parser, argv)
     if args.num_frames <= 0:
         parser.error("--num_frames must be greater than zero")
+    if args.warmup_steps < 0:
+        parser.error("--warmup_steps must be non-negative")
     sys.argv = [sys.argv[0]] + remaining
     return args, remaining
 
@@ -168,6 +176,7 @@ def run(argv: list[str]) -> None:
                             "serialized_synchronized" if args.measure_synchronized_step_breakdown else "host_return"
                         ),
                     },
+                    {"name": "environment_step_warmup_steps", "data": args.warmup_steps},
                     {"name": "presets", "data": ",".join(cfg.presets)},
                 ]
             },
@@ -192,7 +201,9 @@ def run(argv: list[str]) -> None:
         policy = runner.get_inference_policy(device=env.unwrapped.device)
 
         environment_step_timer = stepping.EnvironmentStepTimingRecorder(
-            env, measure_synchronized_step_breakdown=args.measure_synchronized_step_breakdown
+            env,
+            measure_synchronized_step_breakdown=args.measure_synchronized_step_breakdown,
+            warmup_steps=args.warmup_steps,
         )
         with environment_step_timer, BenchmarkMonitor(benchmark, interval=1.0):
             step_times, reward, ep_length, success_rate = stepping.run_play_loop(env, policy, args.num_frames)

scripts/benchmarks/rsl_rl/benchmark_rsl_rl_train.py

@@ -65,6 +65,12 @@ def _parse_args(argv: list[str]):
         action="store_true",
         help="Measure a serialized synchronized simulation and outside-simulation step breakdown.",
     )
+    parser.add_argument(
+        "--warmup_steps",
+        type=int,
+        default=0,
+        help="Exclude the first N env.step() calls from environment-step timing (cold-start). Opt-in; default 0.",
+    )
     parser.add_argument(
         "--benchmark_formatter",
         type=str,
@@ -98,6 +104,8 @@ def _parse_args(argv: list[str]):
     args_cli, remaining_args = setup_preset_cli(parser, argv)
     if args_cli.max_iterations is not None and args_cli.max_iterations <= 0:
         parser.error("--max_iterations must be greater than zero")
+    if args_cli.warmup_steps < 0:
+        parser.error("--warmup_steps must be non-negative")
     enable_cameras_for_video(args_cli)
     sys.argv = [sys.argv[0]] + remaining_args
 
@@ -194,6 +202,7 @@ def run(argv: list[str]) -> None:
                             "serialized_synchronized" if args_cli.measure_synchronized_step_breakdown else "host_return"
                         ),
                     },
+                    {"name": "environment_step_warmup_steps", "data": args_cli.warmup_steps},
                     {"name": "presets", "data": ",".join(cfg.presets)},
                 ]
             },
@@ -232,7 +241,9 @@ def run(argv: list[str]) -> None:
         )
 
         environment_step_timer = stepping.EnvironmentStepTimingRecorder(
-            env, measure_synchronized_step_breakdown=args_cli.measure_synchronized_step_breakdown
+            env,
+            measure_synchronized_step_breakdown=args_cli.measure_synchronized_step_breakdown,
+            warmup_steps=args_cli.warmup_steps,
         )
         with early, environment_step_timer, BenchmarkMonitor(benchmark, interval=1.0):
             runner.learn(num_learning_iterations=agent_cfg.max_iterations, init_at_random_ep_len=True)

scripts/benchmarks/sb3/benchmark_sb3_play.py

@@ -71,6 +71,12 @@ def _parse_args(argv: list[str]):
         action="store_true",
         help="Measure a serialized synchronized simulation and outside-simulation step breakdown.",
     )
+    parser.add_argument(
+        "--warmup_steps",
+        type=int,
+        default=0,
+        help="Exclude the first N env.step() calls from environment-step timing (cold-start). Opt-in; default 0.",
+    )
     parser.add_argument(
         "--benchmark_formatter",
         type=str,
@@ -86,6 +92,8 @@ def _parse_args(argv: list[str]):
     args_cli, remaining_args = setup_preset_cli(parser, argv)
     if args_cli.num_frames <= 0:
         parser.error("--num_frames must be greater than zero")
+    if args_cli.warmup_steps < 0:
+        parser.error("--warmup_steps must be non-negative")
     sys.argv = [sys.argv[0]] + remaining_args
 
     return args_cli, remaining_args
@@ -171,6 +179,7 @@ def run(argv: list[str]) -> None:
                             "serialized_synchronized" if args_cli.measure_synchronized_step_breakdown else "host_return"
                         ),
                     },
+                    {"name": "environment_step_warmup_steps", "data": args_cli.warmup_steps},
                     {"name": "presets", "data": ",".join(cfg.presets)},
                 ]
             },
@@ -222,7 +231,9 @@ def policy(obs):
             return actions
 
         environment_step_timer = stepping.EnvironmentStepTimingRecorder(
-            env, measure_synchronized_step_breakdown=args_cli.measure_synchronized_step_breakdown
+            env,
+            measure_synchronized_step_breakdown=args_cli.measure_synchronized_step_breakdown,
+            warmup_steps=args_cli.warmup_steps,
         )
         with environment_step_timer, BenchmarkMonitor(benchmark, interval=1.0):
             step_times, reward, ep_length, success_rate = stepping.run_play_loop(env, policy, args_cli.num_frames)

scripts/benchmarks/sb3/benchmark_sb3_train.py

@@ -129,6 +129,12 @@ def _parse_args(argv: list[str]):
         action="store_true",
         help="Measure a serialized synchronized simulation and outside-simulation step breakdown.",
     )
+    parser.add_argument(
+        "--warmup_steps",
+        type=int,
+        default=0,
+        help="Exclude the first N env.step() calls from environment-step timing (cold-start). Opt-in; default 0.",
+    )
     parser.add_argument(
         "--benchmark_formatter",
         type=str,
@@ -160,6 +166,8 @@ def _parse_args(argv: list[str]):
     args_cli, remaining_args = setup_preset_cli(parser, argv)
     if args_cli.max_iterations is not None and args_cli.max_iterations <= 0:
         parser.error("--max_iterations must be greater than zero")
+    if args_cli.warmup_steps < 0:
+        parser.error("--warmup_steps must be non-negative")
     enable_cameras_for_video(args_cli)
     sys.argv = [sys.argv[0]] + remaining_args
 
@@ -256,6 +264,7 @@ def run(argv: list[str]) -> None:
                             "serialized_synchronized" if args_cli.measure_synchronized_step_breakdown else "host_return"
                         ),
                     },
+                    {"name": "environment_step_warmup_steps", "data": args_cli.warmup_steps},
                     {"name": "presets", "data": ",".join(cfg.presets)},
                 ]
             },
@@ -312,7 +321,9 @@ def run(argv: list[str]) -> None:
         checkpoint_callback = CheckpointCallback(save_freq=1000, save_path=log_dir, name_prefix="model", verbose=2)
 
         environment_step_timer = stepping.EnvironmentStepTimingRecorder(
-            env, measure_synchronized_step_breakdown=args_cli.measure_synchronized_step_breakdown
+            env,
+            measure_synchronized_step_breakdown=args_cli.measure_synchronized_step_breakdown,
+            warmup_steps=args_cli.warmup_steps,
         )
         with (
             contextlib.suppress(KeyboardInterrupt),

scripts/benchmarks/skrl/benchmark_skrl_play.py

@@ -84,6 +84,12 @@ def _parse_args(argv: list[str]):
         action="store_true",
         help="Measure a serialized synchronized simulation and outside-simulation step breakdown.",
     )
+    parser.add_argument(
+        "--warmup_steps",
+        type=int,
+        default=0,
+        help="Exclude the first N env.step() calls from environment-step timing (cold-start). Opt-in; default 0.",
+    )
     parser.add_argument(
         "--benchmark_formatter",
         type=str,
@@ -99,6 +105,8 @@ def _parse_args(argv: list[str]):
     args_cli, remaining_args = setup_preset_cli(parser, argv)
     if args_cli.num_frames <= 0:
         parser.error("--num_frames must be greater than zero")
+    if args_cli.warmup_steps < 0:
+        parser.error("--warmup_steps must be non-negative")
     sys.argv = [sys.argv[0]] + remaining_args
 
     return args_cli, remaining_args
@@ -200,6 +208,7 @@ def run(argv: list[str]) -> None:
                             "serialized_synchronized" if args_cli.measure_synchronized_step_breakdown else "host_return"
                         ),
                     },
+                    {"name": "environment_step_warmup_steps", "data": args_cli.warmup_steps},
                     {"name": "presets", "data": ",".join(cfg.presets)},
                 ]
             },
@@ -244,7 +253,9 @@ def policy(obs):
             return outputs[-1].get("mean_actions", outputs[0])
 
         environment_step_timer = stepping.EnvironmentStepTimingRecorder(
-            env, measure_synchronized_step_breakdown=args_cli.measure_synchronized_step_breakdown
+            env,
+            measure_synchronized_step_breakdown=args_cli.measure_synchronized_step_breakdown,
+            warmup_steps=args_cli.warmup_steps,
         )
         with environment_step_timer, BenchmarkMonitor(benchmark, interval=1.0):
             step_times, reward, ep_length, success_rate = stepping.run_play_loop(env, policy, args_cli.num_frames)

scripts/benchmarks/skrl/benchmark_skrl_train.py

@@ -166,6 +166,12 @@ def _parse_args(argv: list[str]):
         action="store_true",
         help="Measure a serialized synchronized simulation and outside-simulation step breakdown.",
     )
+    parser.add_argument(
+        "--warmup_steps",
+        type=int,
+        default=0,
+        help="Exclude the first N env.step() calls from environment-step timing (cold-start). Opt-in; default 0.",
+    )
     parser.add_argument(
         "--benchmark_formatter",
         type=str,
@@ -200,6 +206,8 @@ def _parse_args(argv: list[str]):
     args_cli, remaining_args = setup_preset_cli(parser, argv)
     if args_cli.max_iterations is not None and args_cli.max_iterations <= 0:
         parser.error("--max_iterations must be greater than zero")
+    if args_cli.warmup_steps < 0:
+        parser.error("--warmup_steps must be non-negative")
     enable_cameras_for_video(args_cli)
     sys.argv = [sys.argv[0]] + remaining_args
 
@@ -320,6 +328,7 @@ def run(argv: list[str]) -> None:
                             "serialized_synchronized" if args_cli.measure_synchronized_step_breakdown else "host_return"
                         ),
                     },
+                    {"name": "environment_step_warmup_steps", "data": args_cli.warmup_steps},
                     {"name": "presets", "data": ",".join(cfg.presets)},
                 ]
             },
@@ -360,7 +369,9 @@ def _generate_trainer(self, env, cfg, agent):
         bt = runner._trainer
 
         environment_step_timer = stepping.EnvironmentStepTimingRecorder(
-            env, measure_synchronized_step_breakdown=args_cli.measure_synchronized_step_breakdown
+            env,
+            measure_synchronized_step_breakdown=args_cli.measure_synchronized_step_breakdown,
+            warmup_steps=args_cli.warmup_steps,
         )
         with success_context, environment_step_timer, BenchmarkMonitor(benchmark, interval=1.0):
             runner.run()

scripts/benchmarks/test/test_benchmark_smoke.py

@@ -58,6 +58,18 @@ def test_adapters_reject_non_positive_workloads(library: str, workflow: str, arg
     assert exc_info.value.code == 2
 
 
+@pytest.mark.parametrize("library", ["rsl_rl", "rl_games", "skrl", "sb3"])
+@pytest.mark.parametrize("workflow", ["play", "train"])
+def test_adapters_reject_negative_warmup_steps(library: str, workflow: str):
+    """Benchmark adapters reject a negative --warmup_steps (the cold-start exclusion is opt-in, N >= 0)."""
+    module = _load_adapter(library, workflow)
+
+    with pytest.raises(SystemExit) as exc_info:
+        module._parse_args(["--task", _TASK, "--warmup_steps", "-1", "--headless"])
+
+    assert exc_info.value.code == 2
+
+
 @pytest.mark.parametrize(
     (
         "library",

source/isaaclab/changelog.d/antoiner-runtime-benchmark-warmup.minor.rst

@@ -4,6 +4,9 @@ Changed
 * Changed the runtime benchmark to exclude configurable warmup frames and
   report effective throughput over the complete measured interval, and
   report environment-step host-return rates across runtime, play, and training
-  benchmarks. Added an optional serialized synchronized breakdown of time
-  inside and outside simulation calls; the outside-simulation remainder is not
-  classified as overhead.
+  benchmarks. Added an opt-in ``--warmup_steps`` flag to the play and training
+  benchmarks that excludes the first N ``env.step()`` calls (cold start) from
+  the environment-step timing; it defaults to ``0`` (no exclusion). Added an
+  optional serialized synchronized breakdown of time inside and outside
+  simulation calls; the outside-simulation remainder is not classified as
+  overhead.