benchmark_non_rl.py

# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause

"""Script to benchmark the step throughput of a non-RL environment.

Launches a task, steps it for a fixed number of frames, and records per-step
frametimes so downstream tooling can compute effective environment FPS.
"""

"""Launch Isaac Sim Simulator first."""

import argparse
import contextlib
import os
import sys
import time

from isaaclab.app import AppLauncher

from isaaclab_tasks.utils import setup_preset_cli

# add argparse arguments
parser = argparse.ArgumentParser(description="Train an RL agent with RL-Games.")
parser.add_argument("--video", action="store_true", default=False, help="Record videos during training.")
parser.add_argument("--video_length", type=int, default=200, help="Length of the recorded video (in steps).")
parser.add_argument("--video_interval", type=int, default=2000, help="Interval between video recordings (in steps).")
parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.")
parser.add_argument("--task", type=str, default=None, help="Name of the task.")
parser.add_argument("--seed", type=int, default=None, help="Seed used for the environment")
parser.add_argument(
    "--distributed", action="store_true", default=False, help="Run training with multiple GPUs or nodes."
)
parser.add_argument("--num_frames", type=int, default=100, help="Number of environment frames to run benchmark for.")
parser.add_argument(
    "--benchmark_backend",
    type=str,
    default="omniperf",
    choices=[
        "json",
        "osmo",
        "omniperf",
        "summary",
        "LocalLogMetrics",
        "JSONFileMetrics",
        "OsmoKPIFile",
        "OmniPerfKPIFile",
    ],
    help="Benchmarking backend options, defaults omniperf",
)
parser.add_argument("--output_path", type=str, default=".", help="Path to output benchmark results.")

# append AppLauncher cli args
AppLauncher.add_app_launcher_args(parser)
args_cli, hydra_args = setup_preset_cli(parser)
sys.argv = [sys.argv[0]] + hydra_args
if args_cli.video:
    args_cli.enable_cameras = True

sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), "../.."))

from isaaclab.test.benchmark import BaseIsaacLabBenchmark, BenchmarkMonitor
from isaaclab.utils.timer import Timer

from scripts.benchmarks.utils import (
    get_backend_type,
    get_physics_string,
    get_preset_string,
    log_app_start_time,
    log_python_imports_time,
    log_runtime_step_times,
    log_scene_creation_time,
    log_simulation_start_time,
    log_task_start_time,
    log_total_start_time,
)

imports_time_begin = time.perf_counter_ns()

import os
from datetime import datetime

import gymnasium as gym
import numpy as np
import torch

from isaaclab.envs import DirectMARLEnvCfg, DirectRLEnvCfg, ManagerBasedRLEnvCfg
from isaaclab.utils.dict import print_dict

import isaaclab_tasks  # noqa: F401

# PLACEHOLDER: Extension template (do not remove this comment)
with contextlib.suppress(ImportError):
    import isaaclab_tasks_experimental  # noqa: F401
from isaaclab.app import launch_simulation

from isaaclab_tasks.utils import resolve_task_config

imports_time_end = time.perf_counter_ns()


# Create the benchmark
backend_type = get_backend_type(args_cli.benchmark_backend)
benchmark = BaseIsaacLabBenchmark(
    benchmark_name="benchmark_non_rl",
    backend_type=backend_type,
    output_path=args_cli.output_path,
    use_recorders=True,
    frametime_recorders=backend_type in ("summary", "omniperf"),
    output_prefix=f"benchmark_non_rl_{args_cli.task}",
    workflow_metadata={
        "metadata": [
            {"name": "task", "data": args_cli.task},
            {"name": "seed", "data": args_cli.seed},
            {"name": "num_envs", "data": args_cli.num_envs},
            {"name": "num_frames", "data": args_cli.num_frames},
            {"name": "presets", "data": get_preset_string(hydra_args)},
            {"name": "physics", "data": get_physics_string(hydra_args)},
        ]
    },
)


def main(
    env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg,
    app_start_time_begin: int,
    app_start_time_end: int,
):
    """Benchmark without RL in the loop."""

    # override configurations with non-hydra CLI arguments
    env_cfg.scene.num_envs = args_cli.num_envs if args_cli.num_envs is not None else env_cfg.scene.num_envs
    # For distributed training, launch_simulation() already resolved the
    # correct per-rank device; only apply a CLI --device override for
    # non-distributed runs (the default "cuda:0" would clobber the
    # per-rank device otherwise).
    if not args_cli.distributed:
        env_cfg.sim.device = args_cli.device if args_cli.device is not None else env_cfg.sim.device
    env_cfg.seed = args_cli.seed

    # check for invalid combination of CPU device with distributed training
    if args_cli.distributed and args_cli.device is not None and "cpu" in args_cli.device:
        raise ValueError(
            "Distributed training is not supported when using CPU device. "
            "Please use GPU device (e.g., --device cuda) for distributed training."
        )

    # process distributed
    # env_cfg.sim.device is already resolved by launch_simulation().
    world_size = 1
    world_rank = 0
    if args_cli.distributed:
        world_size = int(os.getenv("WORLD_SIZE", 1))
        world_rank = int(os.getenv("RANK", "0"))

    task_startup_time_begin = time.perf_counter_ns()

    # create isaac environment
    env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None)
    # wrap for video recording
    if args_cli.video:
        log_root_path = os.path.abspath(f"benchmark/{args_cli.task}")
        log_dir = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
        video_kwargs = {
            "video_folder": os.path.join(log_root_path, log_dir, "videos"),
            "step_trigger": lambda step: step % args_cli.video_interval == 0,
            "video_length": args_cli.video_length,
            "disable_logger": True,
        }
        print("[INFO] Recording videos during training.")
        print_dict(video_kwargs, nesting=4)
        env = gym.wrappers.RecordVideo(env, **video_kwargs)

    task_startup_time_end = time.perf_counter_ns()

    env.reset()

    # counter for number of frames to run for
    num_frames = 0
    # log frame times
    step_times = []

    # Run with continuous benchmark monitoring
    with BenchmarkMonitor(benchmark, interval=1.0):
        while num_frames < args_cli.num_frames:
            # get upper and lower bounds of action space, sample actions randomly on this interval
            action_high = 1
            action_low = -1
            actions = (action_high - action_low) * torch.rand(
                env.unwrapped.num_envs, env.unwrapped.single_action_space.shape[0], device=env.unwrapped.device
            ) - action_high

            # env stepping
            env_step_time_begin = time.perf_counter_ns()
            _ = env.step(actions)
            end_step_time_end = time.perf_counter_ns()
            step_times.append(end_step_time_end - env_step_time_begin)

            num_frames += 1

    if world_rank == 0:
        # Final update after loop completes
        benchmark.update_manual_recorders()

        # compute stats
        step_times = np.array(step_times) / 1e6  # ns to ms
        fps = 1.0 / (step_times / 1000)
        effective_fps = fps * env.unwrapped.num_envs * world_size

        # prepare step timing dict
        environment_step_times = {
            "Environment step times": step_times.tolist(),
            "Environment step FPS": fps.tolist(),
            "Environment step effective FPS": effective_fps.tolist(),
        }

        log_app_start_time(benchmark, (app_start_time_end - app_start_time_begin) / 1e6)
        log_python_imports_time(benchmark, (imports_time_end - imports_time_begin) / 1e6)
        log_task_start_time(benchmark, (task_startup_time_end - task_startup_time_begin) / 1e6)
        log_scene_creation_time(benchmark, Timer.get_timer_info("scene_creation") * 1000)
        log_simulation_start_time(benchmark, Timer.get_timer_info("simulation_start") * 1000)
        log_total_start_time(benchmark, (task_startup_time_end - app_start_time_begin) / 1e6)
        log_runtime_step_times(benchmark, environment_step_times, compute_stats=True)

        benchmark._finalize_impl()

    # close the simulator
    env.close()


if __name__ == "__main__":
    env_cfg, _agent_cfg = resolve_task_config(args_cli.task, None)

    app_start_time_begin = time.perf_counter_ns()
    with launch_simulation(env_cfg, args_cli):
        app_start_time_end = time.perf_counter_ns()
        main(env_cfg, app_start_time_begin, app_start_time_end)