diff --git a/docs/source/api/lab/isaaclab.utils.rst b/docs/source/api/lab/isaaclab.utils.rst index 9ef89739f79..8a179490556 100644 --- a/docs/source/api/lab/isaaclab.utils.rst +++ b/docs/source/api/lab/isaaclab.utils.rst @@ -20,6 +20,7 @@ timer types warp + mesh .. Rubric:: Functions diff --git a/scripts/benchmarks/aggregated_plots_raycaster.py b/scripts/benchmarks/aggregated_plots_raycaster.py new file mode 100644 index 00000000000..1957b3e4328 --- /dev/null +++ b/scripts/benchmarks/aggregated_plots_raycaster.py @@ -0,0 +1,611 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import matplotlib.pyplot as plt +from matplotlib.ticker import ScalarFormatter +from matplotlib import cm +import argparse +import os +import glob +import numpy as np + +import pandas as pd + +device_name_lookup = { + "NVIDIA RTX 3060": "RTX 3060", + "NVIDIA RTX 3090": "RTX 3090", + "NVIDIA RTX 4090": "RTX 4090", + "NVIDIA RTX 4080": "RTX 4080", + "NVIDIA RTX 4070 Ti": "RTX 4070 Ti", + "NVIDIA RTX 4060 Ti": "RTX 4060 Ti", + "NVIDIA RTX 3080": "RTX 3080", + "NVIDIA RTX 2080 Ti": "RTX 2080 Ti", + "NVIDIA RTX 2080 SUPER": "RTX 2080 SUPER", + "NVIDIA RTX 2080": "RTX 2080", + "NVIDIA GTX 1080 Ti": "GTX 1080 Ti", + "NVIDIA GTX 1080": "GTX 1080", + "NVIDIA GTX 1070": "GTX 1070", + "NVIDIA GTX 1060": "GTX 1060", + "NVIDIA GTX 980 Ti": "GTX 980 Ti", + "NVIDIA GTX 980": "GTX 980", + "NVIDIA RTX 6000 Ada Generation": "RTX 6000 Ada", + "NVIDIA RTX PRO 6000 Blackwell Server Edition": "Blackwell 6000", +} + + +registered_devices = [] +registered_resolutions = {} + +def get_color(device: str, resolution: float) -> str: + """Get a color string based on device and resolution.""" + cmap = cm.get_cmap('tab20c') + base_palette = [cmap(i) for i in range(cmap.N)] + + if device not in registered_devices: + registered_devices.append(device) + registered_resolutions[device] = [] + if resolution is not None and resolution not in registered_resolutions[device]: + registered_resolutions[device].append(resolution) + + base_color_idx = (registered_devices.index(device) * 4) % 20 + if resolution is None: + fine_color_idx = base_color_idx + else: + fine_color_idx = registered_resolutions[device].index(resolution) % 4 + base_color_idx + + return base_palette[fine_color_idx % len(base_palette)] + + +import scienceplots # noqa: F401 +plt.style.use(["science"]) # publication-ready base style + +SMALL_SIZE = 16 +MEDIUM_SIZE = 20 +LEGEND_SIZE = 15 +BIGGER_SIZE = 24 + +plt.rc('font', size=14) # base font size +plt.rc('axes', titlesize=SMALL_SIZE) +plt.rc('axes', labelsize=MEDIUM_SIZE) +plt.rc('xtick', labelsize=SMALL_SIZE) +plt.rc('ytick', labelsize=SMALL_SIZE) +plt.rc('legend', fontsize=LEGEND_SIZE) +plt.rc('figure', titlesize=BIGGER_SIZE) +plt.rc("lines", linewidth=2) + +OUTPUT_DIR = os.environ.get("BENCHMARK_PLOTS_DIR", "scripts/outputs/nvidia/") + + +def _format_axes(ax: plt.Axes, grid: bool = True): + """Apply consistent, publication-ready formatting to axes.""" + if grid: + ax.grid(True, which="major", linestyle="--", linewidth=0.6, alpha=0.6) + ax.grid(True, which="minor", linestyle=":", linewidth=0.4, alpha=0.5) + # Use scalar formatter without offset/scientific on axes unless necessary + ax.xaxis.set_major_formatter(ScalarFormatter(useMathText=True)) + ax.yaxis.set_major_formatter(ScalarFormatter(useMathText=True)) + for spine in ax.spines.values(): + spine.set_linewidth(0.8) + + +def get_dataframe(df_name: str, fields=None, keys=None) -> pd.DataFrame: + """Load CSVs (glob-aware) and optionally aggregate mean/std on requested fields. + + Accepted logical names: + - "single_vs_multi" -> ray_caster_benchmark_single_vs_multi.csv + - "num_assets_reference" -> ray_caster_benchmark_num_assets_reference.csv + - "num_assets" -> ray_caster_benchmark_num_assets.csv + - "num_faces" -> ray_caster_benchmark_num_faces.csv + + If df_name ends with .csv, it will look for that filename; if df_name contains + wildcard characters (*, ?), it will be treated as a glob pattern. Matching + is recursive under the parent folder of OUTPUT_DIR (e.g., scripts/artifacts/**). + + Parameters + - df_name: logical name or CSV filename (or glob pattern) + - fields: optional list of column names to include; if provided, only + existing columns from this list are kept. When None, all columns are kept. + - keys: optional list of column names to group by for aggregation. When + provided (and multiple files match), numeric fields in `fields` that are + not group keys will be aggregated with mean/std. If not provided, data is + concatenated without aggregation. + """ + mapping = { + "single_vs_multi": "ray_caster_benchmark_single_vs_multi.csv", + "num_assets_reference": "ray_caster_benchmark_num_assets_reference.csv", + "num_assets": "ray_caster_benchmark_num_assets.csv", + "num_faces": "ray_caster_benchmark_num_faces.csv", + } + + + pattern_or_name = mapping.get(df_name, df_name) + # Determine search root and pattern + root_dir = os.path.normpath(OUTPUT_DIR) # e.g., scripts/artifacts + print(f"Searching for CSVs matching '{pattern_or_name}' under '{root_dir}'") + if any(ch in pattern_or_name for ch in ["*", "?"]): + pattern = os.path.join(root_dir, "**", pattern_or_name) + print(f"Using glob pattern: {pattern}") + else: + # Search for exact filename under all subfolders + pattern = os.path.join(root_dir, "**", pattern_or_name) + print(f"Using exact filename search with pattern: {pattern}") + + matches = glob.glob(pattern, recursive=True) + # Fallback: direct path inside OUTPUT_DIR + if not matches: + direct = os.path.join(OUTPUT_DIR, pattern_or_name) + if os.path.exists(direct): + matches = [direct] + if not matches: + raise FileNotFoundError(f"No CSV files found for pattern/name: {pattern_or_name} under {root_dir}") + + # Load and optionally column-filter + frames = [] + print(f"Loading {len(matches)} CSV files for pattern/name: {pattern_or_name}") + + for p in matches: + try: + df = pd.read_csv(p) + + # Derived metrics + if fields and ("fps" in fields or "kfps" in fields or "mrays_per_s" in fields): + # Frames per second across all envs (aggregate FPS) + df["fps"] = 1e3 / df["per_step_ms"] * df["num_envs"] # total FPS + df["kfps"] = df["fps"] / 1e3 + if "total_rays" in df.columns: + # Throughput in Mray/s + df["mrays_per_s"] = (df["fps"] * (df["total_rays"] / df["num_envs"])) / 1e6 + + if fields: + keep = [c for c in (fields or []) if c in df.columns] + [c for c in (keys or []) if c in df.columns] + if keep: + df = df[keep] + frames.append(df) + except Exception: + continue + if not frames: + raise RuntimeError(f"Failed to load any CSVs for: {pattern_or_name}") + + df_all = pd.concat(frames, ignore_index=True) + + + + # Determine metrics = numeric fields that are in `fields` but not in `keys` + keys_present = [k for k in (keys or []) if k in df_all.columns] + metrics: list[str] = [] + for c in (fields or []): + if c in keys_present: + continue + try: + if pd.api.types.is_numeric_dtype(df_all[c]): + metrics.append(c) + except Exception: + continue + if not metrics or not keys_present: + return df_all + + group = df_all.groupby(keys_present, dropna=False) + mean_df = group[metrics].mean().reset_index() + std_df = group[metrics].std().reset_index() + # Rename std columns with _std suffix + std_df = std_df.rename(columns={m: f"{m}_std" for m in metrics}) + grouped = pd.merge(mean_df, std_df, on=keys_present, how="left") + # map device names to shorter versions if applicable + if "device" in grouped.columns: + # Map known device names to short labels; keep original where not mapped + grouped["device"] = grouped["device"].map(device_name_lookup).fillna(grouped["device"]) + return grouped + + +# plot1: FPS vs Num Assets. +def plot_num_assets_reference(): + df = get_dataframe( + "num_assets_reference", + fields=["avg_memory", "per_step_ms", "fps", "kfps", "mrays_per_s", "total_rays"], + keys=["reference_meshes", "num_assets", "device", "num_envs", "resolution"], + ) + df = df[df["reference_meshes"] == True] + + fig, axes = plt.subplots(1, 1, figsize=(7.2, 5.61), dpi=600, constrained_layout=True) + for (device, ref), group in df.groupby(["device", "reference_meshes"]): + group = group.sort_values("num_assets") + color = get_color(device, None) + label = f"{device}" + # shade the std deviation + x = group["num_assets"].to_numpy(dtype=float) + y = (group["kfps"].to_numpy(dtype=float) if "kfps" in group.columns else (group["fps"].to_numpy(dtype=float) / 1e3)) + ystd = ( + group["kfps_std"].to_numpy(dtype=float) if "kfps_std" in group.columns else + (group["fps_std"].to_numpy(dtype=float) / 1e3 if "fps_std" in group.columns else np.zeros_like(y)) + ) + axes.fill_between(x, y - ystd, y + ystd, alpha=0.2, color=color) + axes.plot(x, y, marker="o", label=label, color=color) + # Secondary axis as Mray/s + if set(["total_rays", "num_envs"]).issubset(df.columns) and len(df["num_envs"].unique()) == 1: + rays_per_env = (df["total_rays"] / df["num_envs"]).median() + if pd.notnull(rays_per_env) and rays_per_env > 0: + kfps_to_mrays = lambda y: y * rays_per_env / 1e3 + mrays_to_kfps = lambda y: y * 1e3 / rays_per_env + ax_right = axes.secondary_yaxis('right', functions=(kfps_to_mrays, mrays_to_kfps)) + ax_right.set_ylabel(r"Throughput (Rays / s) $\times 10^6$") + axes.set_xlabel("Number of assets") + axes.set_ylabel(r"Throughput (FPS) $\times 10^3$") + #axes.set_title(f"Throughput vs number of assets ({int(df['num_envs'].iloc[0])} envs)") + _format_axes(axes) + axes.legend(frameon=False) + return axes + +# plot2: FPS vs Mesh Complexity. +def plot_mesh_complexity(): + df = get_dataframe( + "num_faces", + fields=["avg_memory", "per_step_ms", "fps", "kfps", "mrays_per_s", "total_rays"], + keys=["num_faces", "device", "num_envs", "resolution"], + ) + + fig, axes = plt.subplots(1, 1, figsize=(7.2, 5.61), dpi=600, constrained_layout=True) + for device, group in df.groupby(["device"]): + group = group.sort_values("num_faces") + color = get_color(device, None) + label = f"{device[0]}" + # shade the std deviation + x = group["num_faces"].to_numpy(dtype=float) + y = (group["kfps"].to_numpy(dtype=float) if "kfps" in group.columns else (group["fps"].to_numpy(dtype=float) / 1e3)) + ystd = ( + group["kfps_std"].to_numpy(dtype=float) if "kfps_std" in group.columns else + (group["fps_std"].to_numpy(dtype=float) / 1e3 if "fps_std" in group.columns else np.zeros_like(y)) + ) + axes.fill_between(x, y - ystd, y + ystd, alpha=0.2, color=color) + axes.plot(x, y, marker="o", label=label, color=color) + # Add right-hand tick labels as Mray/s using a conversion from kFPS + if set(["total_rays", "num_envs"]).issubset(df.columns) and len(df["num_envs"].unique()) == 1: + rays_per_env = (df["total_rays"] / df["num_envs"]).median() + if pd.notnull(rays_per_env) and rays_per_env > 0: + kfps_to_mrays = lambda y: y * rays_per_env / 1e3 + mrays_to_kfps = lambda y: y * 1e3 / rays_per_env + ax_right = axes.secondary_yaxis('right', functions=(kfps_to_mrays, mrays_to_kfps)) + ax_right.set_ylabel(r"Throughput (Rays) $\times 10^6$") + axes.set_xlabel("Mesh complexity (faces)") + axes.set_ylabel(r"Throughput (FPS) $\times 10^3$") + #axes.set_title(f"Throughput vs mesh complexity ({int(df['num_envs'].iloc[0])} envs)") + _format_axes(axes) + axes.legend(frameon=False) + return axes + + +def compare_single_vs_multi_resolution(): + df = get_dataframe( + "single_vs_multi", + fields=["avg_memory", "per_step_ms", "fps", "kfps", "mrays_per_s", "total_rays"], + keys=["device", "num_envs", "resolution", "mode"], + ) + # df = df[df["mode"] == "multi"] + + fig, axes = plt.subplots(1, 1, figsize=(7.2, 5.61), dpi=600, constrained_layout=True) + # build color and line-style maps: color per device, style per resolution + devices = sorted(df["device"].unique()) + style_cycle = ["-", "--", "-.", ":", (0, (3, 1, 1, 1)), (0, (5, 2, 1, 2))] + style_map = {} + for dev in devices: + res_list = sorted(df[df["device"] == dev]["resolution"].unique()) + for idx, r in enumerate(res_list): + style_map[(dev, r)] = style_cycle[idx % len(style_cycle)] + + idx = 0 + # group by device and resolution + for (device, res, mode), group in df.groupby(["device", "resolution", "mode"]): + if mode != "multi": + continue + color = get_color(device, res) + group = group.sort_values("num_envs") + label = f"{device} - ({(5.0 / res)**1:.0f}$\\times${(5.0 / res)**1:.0f})" # - {mode}" + # shade the std deviation with device color + x = group["num_envs"].to_numpy(dtype=float) + y = (group["mrays_per_s"].to_numpy(dtype=float) if "mrays_per_s" in group.columns else (group["fps"].to_numpy(dtype=float) / 1e3)) + ystd = ( + group["mrays_per_s_std"].to_numpy(dtype=float) if "mrays_per_s_std" in group.columns else + (group["fps_std"].to_numpy(dtype=float) / 1e3 if "fps_std" in group.columns else np.zeros_like(y)) + ) + axes.fill_between(list(x), list(y - ystd), list(y + ystd), alpha=0.2, color=color) + axes.plot(list(x), list(y), marker="o", label=label, linestyle=style_map.get((device, res), "-"), color=color) + axes.set_xlabel("Number of environments") + axes.set_ylabel(r"Throughput (Rays / s) $\times 10^6$") + #axes.set_title("Throughput vs number of environments") + _format_axes(axes) + axes.legend(frameon=False, ncol=1, loc="best") + return axes + + +def compare_single_vs_multi_fps(): + df = get_dataframe( + "single_vs_multi", + fields=["avg_memory", "per_step_ms", "fps", "kfps", "mrays_per_s", "total_rays"], + keys=["device", "num_envs", "resolution", "mode"], + ) + # df = df[df["mode"] == "multi"] + + fig, axes = plt.subplots(1, 1, figsize=(7.2, 5.61), dpi=600, constrained_layout=True) + # build color and line-style maps: color per device, style per resolution + devices = sorted(df["device"].unique()) + style_cycle = ["-", "--", "-.", ":", (0, (3, 1, 1, 1)), (0, (5, 2, 1, 2))] + style_map = {} + for dev in devices: + res_list = sorted(df[df["device"] == dev]["resolution"].unique()) + for idx, r in enumerate(res_list): + style_map[(dev, r)] = style_cycle[idx % len(style_cycle)] + + idx = 0 + # group by device and resolution + for (device, res, mode), group in df.groupby(["device", "resolution", "mode"]): + if mode != "multi": + continue + color = get_color(device, res) + group = group.sort_values("num_envs") + label = f"{device} - ({(5.0 / res)**1:.0f}$\\times${(5.0 / res)**1:.0f})" # - {mode}" + # shade the std deviation with device color + x = group["num_envs"].to_numpy(dtype=float) + y = (group["kfps"].to_numpy(dtype=float) if "kfps" in group.columns else (group["fps"].to_numpy(dtype=float) / 1e3)) + ystd = ( + group["kfps_std"].to_numpy(dtype=float) if "kfps_std" in group.columns else + (group["fps_std"].to_numpy(dtype=float) / 1e3 if "fps_std" in group.columns else np.zeros_like(y)) + ) + axes.fill_between(list(x), list(y - ystd), list(y + ystd), alpha=0.2, color=color) + axes.plot(list(x), list(y), marker="o", label=label, linestyle=style_map.get((device, res), "-"), color=color) + axes.set_xlabel("Number of environments") + axes.set_ylabel(r"Throughput (FPS) $\times 10^3$") + #axes.set_title("Throughput vs number of environments") + _format_axes(axes) + axes.legend(frameon=False, ncol=1, loc="best") + return axes + +def compare_memory_consumption(): + df = get_dataframe( + "single_vs_multi", + fields=["avg_memory", "per_step_ms", "fps", "kfps", "mrays_per_s", "total_rays"], + keys=["device", "num_envs", "resolution", "mode"], + ) + df = df[df["mode"] == "multi"] + fig, axes = plt.subplots(1, 1, figsize=(7.2, 5.61), dpi=600, constrained_layout=True) + # build color and line-style maps: color per device, style per resolution + devices = sorted(df["device"].unique()) + style_cycle = ["-", "--", "-.", ":", (0, (3, 1, 1, 1)), (0, (5, 2, 1, 2))] + style_map = {} + for dev in devices: + res_list = sorted(df[df["device"] == dev]["resolution"].unique()) + for idx, r in enumerate(res_list): + style_map[(dev, r)] = style_cycle[idx % len(style_cycle)] + + # group by device and resolution + for (device, res), group in df.groupby(["device", "resolution"]): + label = f"{device} - ({(5.0 / res)**1:.0f}$\\times${(5.0 / res)**1:.0f}) rays" + color = get_color(device, res) + # shade the std deviation with device color + x = group["num_envs"].to_numpy(dtype=float) + # normalize each series by its own minimum + y_raw = group["avg_memory"].to_numpy(dtype=float) + # use nanmin to be safe if NaNs are present; fall back to 0 if all NaN + y_min = np.nanmin(y_raw) if np.any(~np.isnan(y_raw)) else 0.0 + y = y_raw - y_min + ystd = group["avg_memory_std"].to_numpy(dtype=float) if "avg_memory_std" in group.columns else np.zeros_like(y) + axes.fill_between(x, y - ystd, y + ystd, alpha=0.2, color=color) + axes.plot(x, y, marker="o", label=label, linestyle=style_map.get((device, res), "-"), color=color) + axes.set_xlabel("Number of environments") + axes.set_ylabel(r"VRAM usage $\Delta$ (MB)") + # saxes.set_title("VRAM usage vs number of environments") + _format_axes(axes) + axes.legend(frameon=False, loc="best") + return axes + +if __name__ == "__main__": + os.makedirs(OUTPUT_DIR, exist_ok=True) + try: + axes = plot_num_assets_reference() + plt.savefig(os.path.join(OUTPUT_DIR, "num_assets_reference.png"), dpi=800) + print(f"Saved plot to {os.path.join(OUTPUT_DIR, 'num_assets_reference.png')}") + except Exception as e: + print(f"Failed to plot num_assets_reference: {e}") + + try: + axes = plot_mesh_complexity() + plt.savefig(os.path.join(OUTPUT_DIR, "mesh_complexity.png"), dpi=800) + print(f"Saved plot to {os.path.join(OUTPUT_DIR, 'mesh_complexity.png')}") + except Exception as e: + print(f"Failed to plot mesh_complexity: {e}") + + try: + axes = compare_single_vs_multi_resolution() + plt.savefig(os.path.join(OUTPUT_DIR, "single_vs_multi_resolution.png"), dpi=800) + print(f"Saved plot to {os.path.join(OUTPUT_DIR, 'single_vs_multi_resolution.png')}") + except Exception as e: + print(f"Failed to plot single_vs_multi_resolution: {e}") + + try: + axes = compare_single_vs_multi_fps() + plt.savefig(os.path.join(OUTPUT_DIR, "single_vs_multi_fps.png"), dpi=800) + print(f"Saved plot to {os.path.join(OUTPUT_DIR, 'single_vs_multi_fps.png')}") + except Exception as e: + print(f"Failed to plot single_vs_multi_fps: {e}") + + try: + axes = compare_memory_consumption() + plt.savefig(os.path.join(OUTPUT_DIR, "memory_consumption.png"), dpi=800) + print(f"Saved plot to {os.path.join(OUTPUT_DIR, 'memory_consumption.png')}") + except Exception as e: + print(f"Failed to plot memory_consumption: {e}") + # Show all plots (optional; comment out for headless) + plt.show() + + +# # Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# # All rights reserved. +# # +# # SPDX-License-Identifier: BSD-3-Clause + +# import os + +# import pandas as pd +# import plotly.express as px + +# import wandb + + +# def compare_single_vs_multi(): +# csv_path = "outputs/benchmarks/ray_caster_benchmark_single_vs_multi.csv" +# df = pd.read_csv(csv_path) + +# artifact = wandb.Artifact("ray_caster_benchmark_single_vs_multi", type="dataset") +# artifact.add_file(csv_path) +# wandb.log_artifact(artifact) + +# df["resolution"] = df["resolution"].astype(float) +# df["num_envs"] = df["num_envs"].astype(int) +# df["avg_memory"] = df["avg_memory"].astype(float) +# df["time_per_ray_us"] = df["per_step_ms"] * 1000.0 / df["total_rays"] +# df["rays_per_env"] = df["total_rays"] / df["num_envs"] +# df["fps"] = 1.0 / (df["per_step_ms"] * 1e-3) +# df["rays_per_us"] = 1.0 / df["time_per_ray_us"] + +# df["resolution_label"] = df.apply(lambda r: f"{r['mode']} - {(5.0 / r['resolution'])**2:.2f} rays", axis=1) + +# fig_fps = px.line( +# df, +# x="num_envs", +# y="fps", +# color="resolution_label", +# markers=True, +# title="FPS vs Num Envs", +# ) + +# fig_rays = px.line( +# df, +# x="num_envs", +# y="rays_per_us", +# color="resolution_label", +# markers=True, +# title="Rays per µs vs Num Envs", +# ) + +# fig_mem = px.line( +# df, +# x="num_envs", +# y="avg_memory", +# color="resolution_label", +# markers=True, +# title="Average VRAM Usage (MB) vs Num Envs", +# ) + +# wandb.log({ +# "FPS vs Num Envs": fig_fps, +# "Rays per µs vs Num Envs": fig_rays, +# "VRAM vs Num Envs": fig_mem, +# "Single vs Multi Table": wandb.Table(dataframe=df), +# }) + + +# def compare_num_assets_vram_cache(): +# csv_path = "outputs/benchmarks/ray_caster_benchmark_num_assets_reference.csv" +# df = pd.read_csv(csv_path) + +# artifact = wandb.Artifact("ray_caster_benchmark_num_assets_reference", type="dataset") +# artifact.add_file(csv_path) +# wandb.log_artifact(artifact) + +# df["num_assets"] = df["num_assets"].astype(int) +# df["avg_memory"] = df["avg_memory"].astype(float) +# df["fps"] = 1.0 / (df["per_step_ms"] * 1e-3) +# df["reference_meshes"] = df["reference_meshes"].astype(bool) + +# fig = px.line( +# df, +# x="num_assets", +# y="avg_memory", +# color="reference_meshes", +# markers=True, +# title="Average VRAM Usage (MB) vs Num Assets", +# ) + +# wandb.log({ +# "Num Assets (VRAM Cache)": fig, +# "Num Assets VRAM Table": wandb.Table(dataframe=df), +# }) + + +# def compare_num_assets(): +# csv_path = "outputs/benchmarks/ray_caster_benchmark_num_assets.csv" +# df = pd.read_csv(csv_path) + +# artifact = wandb.Artifact("ray_caster_benchmark_num_assets", type="dataset") +# artifact.add_file(csv_path) +# wandb.log_artifact(artifact) + +# df["num_assets"] = df["num_assets"].astype(int) +# df["avg_memory"] = df["avg_memory"].astype(float) +# df["fps"] = 1.0 / (df["per_step_ms"] * 1e-3) + +# fig_fps = px.line( +# df, +# x="num_assets", +# y="fps", +# markers=True, +# title="FPS vs Num Assets", +# ) + +# fig_mem = px.line( +# df, +# x="num_assets", +# y="avg_memory", +# markers=True, +# title="Average VRAM Usage (MB) vs Num Assets", +# ) + +# wandb.log({ +# "FPS vs Num Assets": fig_fps, +# "VRAM vs Num Assets": fig_mem, +# "Num Assets Table": wandb.Table(dataframe=df), +# }) + + +# def compare_num_faces(): +# csv_path = "outputs/benchmarks/ray_caster_benchmark_num_faces.csv" +# df = pd.read_csv(csv_path) + +# artifact = wandb.Artifact("ray_caster_benchmark_num_faces", type="dataset") +# artifact.add_file(csv_path) +# wandb.log_artifact(artifact) + +# df["num_faces"] = df["num_faces"].astype(int) +# df["avg_memory"] = df["avg_memory"].astype(float) +# df["fps"] = 1.0 / (df["per_step_ms"] * 1e-3) + +# fig_fps = px.line( +# df, +# x="num_faces", +# y="fps", +# markers=True, +# title="FPS vs Num Faces", +# ) + +# fig_mem = px.line( +# df, +# x="num_faces", +# y="avg_memory", +# markers=True, +# title="Average VRAM Usage (MB) vs Num Faces", +# ) + +# wandb.log({ +# "FPS vs Num Faces": fig_fps, +# "VRAM vs Num Faces": fig_mem, +# "Num Faces Table": wandb.Table(dataframe=df), +# }) + + +# if __name__ == "__main__": +# wandb.init(project="raycast-benchmarks", job_type="analysis", entity=os.getenv("WANDB_ENTITY", None)) +# compare_single_vs_multi() +# compare_num_assets_vram_cache() +# compare_num_assets() +# compare_num_faces() +# wandb.finish() diff --git a/scripts/benchmarks/benchmark_camera_throughput.py b/scripts/benchmarks/benchmark_camera_throughput.py new file mode 100644 index 00000000000..1ef18b6e972 --- /dev/null +++ b/scripts/benchmarks/benchmark_camera_throughput.py @@ -0,0 +1,334 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Benchmark throughput of different camera implementations. + +This script benchmarks per-step time while varying: +- camera implementation: standard camera, tiled camera, warp ray-caster camera +- image resolutions (height x width) +- number of environments + +Sensors are added to the scene config before `InteractiveScene` is constructed. +Each benchmark run initializes a fresh simulation and scene and tears it down. + +Examples: + + - Benchmark all camera types across resolutions: + ./isaaclab.sh -p scripts/benchmarks/benchmark_camera_throughput.py \\ + --num_envs 256 512 \\ + --resolutions 240x320,480x640 --steps 200 --warmup 20 --headless + + - Only standard camera at 720p: + ./isaaclab.sh -p scripts/benchmarks/benchmark_camera_throughput.py \\ + --num_envs 256 --resolutions 720x1280 --steps 200 --warmup 20 --headless +""" + +"""Launch Isaac Sim Simulator first.""" + +import argparse +import os +import time + +# local imports +from local_utils import dataframe_to_markdown + +from isaaclab.app import AppLauncher + +parser = argparse.ArgumentParser(description="Benchmark throughput of different camera implementations.") +parser.add_argument( + "--num_envs", + type=int, + nargs="+", + default=[12, 24, 48], # [256, 512, 1024], + help="List of environment counts to benchmark (e.g., 256 512 1024).", +) +parser.add_argument( + "--usd_camera", + action="store_true", + default=False, + help="Whether to benchmark the USD camera.", +) +parser.add_argument( + "--tiled_camera", + action="store_true", + default=False, + help="Whether to benchmark the tiled camera.", +) +parser.add_argument( + "--ray_caster_camera", + action="store_true", + default=False, + help="Whether to benchmark the ray caster camera.", +) +parser.add_argument( + "--resolutions", + type=str, + default="240x320,480x640", + help="Comma-separated list of HxW resolutions, e.g., 240x320,480x640", +) +parser.add_argument( + "--data_type", + type=str, + default="distance_to_image_plane", + help="Data type, e.g., distance_to_image_plane,rgb", +) +parser.add_argument("--steps", type=int, default=500, help="Steps per run to time.") +parser.add_argument("--warmup", type=int, default=50, help="Warmup steps per run before timing.") + +# Append AppLauncher CLI args and parse +AppLauncher.add_app_launcher_args(parser) +args_cli, _ = parser.parse_known_args() +if args_cli.tiled_camera or args_cli.usd_camera: + args_cli.enable_cameras = True +args_cli.headless = True + +# launch omniverse app +app_launcher = AppLauncher(args_cli) +simulation_app = app_launcher.app + +"""Rest everything follows.""" +import torch +import pandas as pd + +import isaacsim.core.utils.stage as stage_utils +from isaacsim.core.simulation_manager import SimulationManager + +import isaaclab.sim as sim_utils +from isaaclab.assets import ArticulationCfg, AssetBaseCfg, RigidObjectCfg +from isaaclab.scene import InteractiveScene, InteractiveSceneCfg +from isaaclab.sensors import CameraCfg, MultiMeshRayCasterCameraCfg, TiledCameraCfg, patterns, Camera, TiledCamera, MultiMeshRayCasterCamera +from isaaclab.sim import SimulationContext +from isaaclab.utils import configclass +from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR + +# Robot config to attach sensors under a valid prim +from isaaclab_assets.robots.anymal import ANYMAL_D_CFG # isort: skip + + +def _parse_resolutions(res_str: str) -> list[tuple[int, int]]: + resolutions: list[tuple[int, int]] = [] + for token in [s for s in res_str.split(",") if s]: + h, w = token.lower().split("x") + resolutions.append((int(h), int(w))) + print("[INFO]: Resolutions: ", resolutions) + return resolutions + + +@configclass +class CameraBenchmarkSceneCfg(InteractiveSceneCfg): + """Scene config with ground, light, robot, and one camera sensor per env.""" + + ground = AssetBaseCfg( + prim_path="/World/ground", + spawn=sim_utils.UsdFileCfg(usd_path=f"{ISAAC_NUCLEUS_DIR}/Environments/Terrains/rough_plane.usd"), + ) + light = AssetBaseCfg( + prim_path="/World/Light", + spawn=sim_utils.DomeLightCfg(intensity=3000.0, color=(0.75, 0.75, 0.75)), + ) + robot: ArticulationCfg = ANYMAL_D_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") # type: ignore[attr-defined] + + # one cube per environment (optional target for ray-caster camera) + cube: RigidObjectCfg = RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/cube", + spawn=sim_utils.CuboidCfg( + size=(0.2, 0.2, 0.2), + rigid_props=sim_utils.RigidBodyPropertiesCfg(max_depenetration_velocity=1.0, disable_gravity=True), + mass_props=sim_utils.MassPropertiesCfg(mass=1.0), + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.5, 0.0, 0.0)), + ), + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, 1.0)), + ) + + usd_camera: CameraCfg | None = None + tiled_camera: TiledCameraCfg | None = None + ray_caster_camera: MultiMeshRayCasterCameraCfg | None = None + + +def _make_scene_cfg_usd(num_envs: int, height: int, width: int, data_types: list[str], debug_vis: bool) -> CameraBenchmarkSceneCfg: + scene_cfg = CameraBenchmarkSceneCfg(num_envs=num_envs, env_spacing=2.0) + scene_cfg.usd_camera = CameraCfg( + prim_path="{ENV_REGEX_NS}/Camera", + height=height, + width=width, + data_types=data_types, + spawn=sim_utils.PinholeCameraCfg( + focal_length=24.0, focus_distance=400.0, horizontal_aperture=20.955, clipping_range=(0.1, 1e4) + ), + debug_vis=debug_vis, + ) + return scene_cfg + + +def _make_scene_cfg_tiled(num_envs: int, height: int, width: int, data_types: list[str], debug_vis: bool) -> CameraBenchmarkSceneCfg: + scene_cfg = CameraBenchmarkSceneCfg(num_envs=num_envs, env_spacing=2.0) + scene_cfg.tiled_camera = TiledCameraCfg( + prim_path="{ENV_REGEX_NS}/TiledCamera", + height=height, + width=width, + data_types=data_types, + spawn=sim_utils.PinholeCameraCfg( + focal_length=24.0, focus_distance=400.0, horizontal_aperture=20.955, clipping_range=(0.1, 1e4) + ), + debug_vis=debug_vis, + ) + return scene_cfg + + +def _make_scene_cfg_ray_caster(num_envs: int, height: int, width: int, data_types: list[str], debug_vis: bool) -> CameraBenchmarkSceneCfg: + scene_cfg = CameraBenchmarkSceneCfg(num_envs=num_envs, env_spacing=2.0) + scene_cfg.ray_caster_camera = MultiMeshRayCasterCameraCfg( + prim_path="{ENV_REGEX_NS}/Robot/base", # attach to existing prim + mesh_prim_paths=["/World/ground", "/World/envs/env_.*/cube"], + pattern_cfg=patterns.PinholeCameraPatternCfg( + focal_length=24.0, horizontal_aperture=20.955, height=height, width=width + ), + data_types=data_types, + debug_vis=debug_vis, + ) + return scene_cfg + + +def _setup_scene(scene_cfg: CameraBenchmarkSceneCfg) -> tuple[SimulationContext, InteractiveScene, float]: + # Create a new stage to avoid residue across runs + stage_utils.create_new_stage() + sim_cfg = sim_utils.SimulationCfg(device=args_cli.device) + sim = SimulationContext(sim_cfg) + sim.set_camera_view((2.5, 0.0, 4.0), (0.0, 0.0, 2.0)) + setup_time_begin = time.perf_counter_ns() + scene = InteractiveScene(scene_cfg) + setup_time_end = time.perf_counter_ns() + print(f"[INFO]: Scene creation time: {(setup_time_end - setup_time_begin) / 1e6:.2f} ms") + reset_time_begin = time.perf_counter_ns() + sim.reset() + reset_time_end = time.perf_counter_ns() + print(f"[INFO]: Sim start time: {(reset_time_end - reset_time_begin) / 1e6:.2f} ms") + return sim, scene, sim.get_physics_dt() + + +def _run_benchmark(scene_cfg: CameraBenchmarkSceneCfg, sensor_name: str): + sim, scene, sim_dt = _setup_scene(scene_cfg) + sensor: Camera | TiledCamera | MultiMeshRayCasterCamera = scene[sensor_name] + # Warmup + for _ in range(args_cli.warmup): + sim.step() + sensor.update(dt=sim_dt, force_recompute=True) + + used_memory = 0.0 + + # Timing + t0 = time.perf_counter_ns() + for _ in range(args_cli.steps): + sim.step() + sensor.update(dt=sim_dt, force_recompute=True) + free, total = torch.cuda.mem_get_info(args_cli.device) + used_memory += (total - free) / 1024**2 # Convert to MB + t1 = time.perf_counter_ns() + per_step_ms = (t1 - t0) / args_cli.steps / 1e6 + avg_memory = used_memory / args_cli.steps + # Cleanup + # stop simulation + # note: cannot use self.sim.stop() since it does one render step after stopping!! This doesn't make sense :( + sim._timeline.stop() + # clear the stage + sim.clear_all_callbacks() + sim.clear_instance() + SimulationManager._simulation_manager_interface.reset() + SimulationManager._callbacks.clear() + + return { + "num_envs": scene.num_envs, + "resolution": sensor.image_shape, + "per_step_ms": float(per_step_ms), + "avg_memory": float(avg_memory), + } + + +def main(): + """Main function.""" + # Prepare benchmark + resolutions = _parse_resolutions(args_cli.resolutions) + cameras = [] + if args_cli.usd_camera: + cameras.append("usd_camera") + if args_cli.tiled_camera: + cameras.append("tiled_camera") + if args_cli.ray_caster_camera: + cameras.append("ray_caster_camera") + data_types = [args_cli.data_type] + device_name = ( + torch.cuda.get_device_name(torch.cuda.current_device()) if torch.cuda.is_available() else platform.processor() + ) + + # BENCHMARK 1 - Compare Depth Camera + print(f"=== Benchmarking {args_cli.data_type} CAMERA ===") + results: list[dict[str, object]] = [] + + for idx, num_envs in enumerate(args_cli.num_envs): + print(f"\n[INFO]: Benchmarking with {num_envs} envs. {idx + 1} / {len(args_cli.num_envs)}") + for resolution in resolutions: + + for camera in cameras: + # USD Camera + if camera == "usd_camera": + single_scene_cfg = _make_scene_cfg_usd( + num_envs=num_envs, + height=resolution[0], + width=resolution[1], + data_types=data_types, + debug_vis=not args_cli.headless, + ) + result = _run_benchmark(single_scene_cfg, "usd_camera") + + # Tiled Camera + elif camera == "tiled_camera": + single_scene_cfg = _make_scene_cfg_tiled( + num_envs=num_envs, + height=resolution[0], + width=resolution[1], + data_types=data_types, + debug_vis=not args_cli.headless, + ) + result = _run_benchmark(single_scene_cfg, "tiled_camera") + + # Multi-Mesh RayCaster Camera + elif camera == "ray_caster_camera": + + if args_cli.data_type == "rgb": + continue + + single_scene_cfg = _make_scene_cfg_ray_caster( + num_envs=num_envs, + height=resolution[0], + width=resolution[1], + data_types=data_types, + debug_vis=not args_cli.headless, + ) + result = _run_benchmark(single_scene_cfg, "ray_caster_camera") + + result["num_envs"] = num_envs + result["resolution"] = resolution + result["mode"] = camera + result["data_types"] = data_types + results.append(result) + del single_scene_cfg + + df_camera = pd.DataFrame(results) + df_camera["device"] = device_name + os.makedirs("outputs/benchmarks", exist_ok=True) + df_camera.to_csv(f"outputs/benchmarks/camera_{args_cli.data_type}_USD_{args_cli.usd_camera}_Tiled_{args_cli.tiled_camera}_RayCaster_{args_cli.ray_caster_camera}_Resolution_{args_cli.resolutions}.csv", index=False) + + # Create .md file with all three tables + for df, title in zip( + [df_camera], [args_cli.data_type] + ): + with open(f"outputs/benchmarks/camera_benchmark_USD_{args_cli.usd_camera}_Tiled_{args_cli.tiled_camera}_RayCaster_{args_cli.ray_caster_camera}_Resolution_{args_cli.resolutions}_{title}.md", "w") as f: + f.write(f"# {title}\n\n") + f.write(dataframe_to_markdown(df, floatfmt=".3f")) + f.write("\n\n") + +if __name__ == "__main__": + main() + simulation_app.close() diff --git a/scripts/benchmarks/benchmark_ray_caster.py b/scripts/benchmarks/benchmark_ray_caster.py new file mode 100644 index 00000000000..47e6e30755c --- /dev/null +++ b/scripts/benchmarks/benchmark_ray_caster.py @@ -0,0 +1,397 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Benchmark ray-caster sensors (single-mesh and multi-mesh). + +This script creates a simple scene with: +- a ground plane +- 0 to N moving spheres as rigid bodies + +It then runs three different benchmarks: +1. Single vs Multi mesh ray caster against ground only +2. Multi-mesh ray caster against ground + various number of spheres +3. Multi-mesh ray caster against spheres with different numbers of vertices +""" + +"""Launch Isaac Sim Simulator first.""" + +import argparse +import numpy as np +import os +import platform +import time +import torch + +import pandas as pd + +from local_utils import dataframe_to_markdown + +from isaaclab.app import AppLauncher + +# add argparse arguments +parser = argparse.ArgumentParser(description="Benchmark ray caster sensors.") +parser.add_argument("--steps", type=int, default=2000, help="Steps per resolution for timing.") +parser.add_argument("--warmup", type=int, default=50, help="Warmup steps before timing.") + +# Num assets for benchmarking memory usage with and without caching. +NUM_ASSETS_MEMORY = [1, 2, 4, 8, 16, 32] +# Num assets for benchmarking scaling performance of multi-mesh ray caster. +NUM_ASSETS = [0, 1, 2, 4, 8, 16, 32] +# Num envs for benchmarking single vs multi mesh ray caster. +NUM_ENVS = [32, 64, 128, 256, 512, 1024, 2048, 4096] +# Num subdivisions for benchmarking mesh complexity. +MESH_SUBDIVISIONS = [0, 1, 2, 3, 4, 5] +# Different ray caster resolutions to benchmark. Num rays will be (5 / res)^2, e.g. 625, 2500, 10000, 11111 +RESOLUTIONS: list[float] = [0.2, 0.1, 0.05, 0.015] + +# Output directory for benchmark artifacts (can be overridden via env var BENCHMARK_OUTPUT_DIR) +timestamp = time.strftime("%Y-%m-%d_%H-%M-%S", time.localtime()) +OUTPUT_DIR = "outputs/benchmarks/" + timestamp + +# # TINY for debugging +# NUM_ASSETS_MEMORY = [1, 2] # Num assets for benchmarking memory usage with and without caching. +# NUM_ASSETS = [0, 1] # Num assets for benchmarking scaling performance. of multi-mesh ray caster. +# NUM_ENVS = [32, 64] # Num envs for benchmarking single vs multi mesh ray caster. +# MESH_SUBDIVISIONS = [0, 1] # Num subdivisions for benchmarking mesh complexity. +# RESOLUTIONS: list[float] = [0.2, 0.1] # Different ray caster resolutions to benchmark. Num rays will be (5 / res)^2, e.g. 625, 2500, 10000, 11111 + +# append AppLauncher cli args +AppLauncher.add_app_launcher_args(parser) +# parse the arguments +args_cli, _ = parser.parse_known_args() +args_cli.headless = True + +# launch omniverse app +app_launcher = AppLauncher(args_cli) +simulation_app = app_launcher.app + +"""Rest everything follows.""" + +import isaacsim.core.utils.stage as stage_utils +from isaacsim.core.simulation_manager import SimulationManager + +import isaaclab.sim as sim_utils +from isaaclab.assets import AssetBaseCfg, RigidObjectCfg +from isaaclab.scene import InteractiveScene, InteractiveSceneCfg +from isaaclab.sensors.ray_caster import MultiMeshRayCasterCfg, RayCaster, RayCasterCfg, patterns +from isaaclab.sim import SimulationContext +from isaaclab.utils import configclass +from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR +from isaaclab.utils.mesh import _MESH_CONVERTERS_CALLBACKS, _create_sphere_trimesh + + +@configclass +class RayCasterBenchmarkSceneCfg(InteractiveSceneCfg): + """Scene config with ground, robot, and optional spheres per env.""" + + # ground plane (rough) + ground = AssetBaseCfg( + prim_path="/World/ground", + spawn=sim_utils.UsdFileCfg(usd_path=f"{ISAAC_NUCLEUS_DIR}/Environments/Terrains/rough_plane.usd"), + ) + + # lights + dome_light = AssetBaseCfg( + prim_path="/World/Light", spawn=sim_utils.DomeLightCfg(intensity=3000.0, color=(0.75, 0.75, 0.75)) + ) + + robot: RigidObjectCfg = RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/ray_caster_origin", + spawn=sim_utils.SphereCfg( + radius=0.1, + rigid_props=sim_utils.RigidBodyPropertiesCfg(kinematic_enabled=True), + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.0, 0.6, 0.0)), + ), + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, 4.0)), + ) + + # spheres collection (optionally set at runtime) + spheres: RigidObjectCfg = RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/sphere", + spawn=sim_utils.SphereCfg( + radius=0.2, + rigid_props=sim_utils.RigidBodyPropertiesCfg(max_depenetration_velocity=1.0, disable_gravity=True), + mass_props=sim_utils.MassPropertiesCfg(mass=100.0), + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.5, 0.0, 0.0)), + collision_props=sim_utils.CollisionPropertiesCfg(), + ), + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, 1.0)), + ) + + # sensors (set at runtime) + height_scanner: RayCasterCfg | None = None + height_scanner_multi: MultiMeshRayCasterCfg | None = None + + +def _make_scene_cfg_single(num_envs: int, resolution: float, debug_vis: bool) -> RayCasterBenchmarkSceneCfg: + scene_cfg = RayCasterBenchmarkSceneCfg(num_envs=num_envs, env_spacing=2.0) + scene_cfg.height_scanner = RayCasterCfg( + prim_path="{ENV_REGEX_NS}/ray_caster_origin", + mesh_prim_paths=["/World/ground"], + pattern_cfg=patterns.GridPatternCfg(resolution=resolution, size=(5.0, 5.0)), + ray_alignment="world", + offset=RayCasterCfg.OffsetCfg(pos=(0.0, 0.0, 20.0)), + debug_vis=debug_vis, + ) + return scene_cfg + + +def _make_scene_cfg_multi( + num_envs: int, + resolution: float, + debug_vis: bool, + track_mesh_transforms: bool, + num_assets: int = 1, + reference_meshes: bool = True, +) -> RayCasterBenchmarkSceneCfg: + scene_cfg = RayCasterBenchmarkSceneCfg(num_envs=num_envs, env_spacing=2.0) + + obj_cfg = scene_cfg.spheres + if track_mesh_transforms: + # Enable gravity + obj_cfg.spawn.rigid_props.disable_gravity = False + + for i in range(num_assets): + new_obj_cfg = obj_cfg.replace(prim_path=f"{{ENV_REGEX_NS}}/sphere_{i}") + ratio = i / num_assets + new_obj_cfg.init_state.pos = (ratio - 0.5, ratio - 0.5, 1.0) + setattr(scene_cfg, f"sphere_{i}", new_obj_cfg) + del scene_cfg.spheres + + scene_cfg.height_scanner_multi = MultiMeshRayCasterCfg( + prim_path="{ENV_REGEX_NS}/ray_caster_origin", + mesh_prim_paths=["/World/ground"] + + [ + MultiMeshRayCasterCfg.RaycastTargetCfg( + target_prim_expr=f"/World/envs/env_.*/sphere_{i}", track_mesh_transforms=track_mesh_transforms + ) + for i in range(num_assets) + ], + pattern_cfg=patterns.GridPatternCfg(resolution=resolution, size=(5.0, 5.0)), + offset=RayCasterCfg.OffsetCfg(pos=(0.0, 0.0, 20.0)), + debug_vis=debug_vis, + ray_alignment="world", + reference_meshes=reference_meshes, + ) + return scene_cfg + + +def _setup_scene(scene_cfg: RayCasterBenchmarkSceneCfg) -> tuple[SimulationContext, InteractiveScene, float]: + # Create a new stage + stage_utils.create_new_stage() + # New simulation per run + sim_cfg = sim_utils.SimulationCfg(device=args_cli.device) + sim = SimulationContext(sim_cfg) + sim.set_camera_view((2.5, 0.0, 4.0), (0.0, 0.0, 2.0)) + # Create scene with sensors + setup_time_begin = time.perf_counter_ns() + scene = InteractiveScene(scene_cfg) + setup_time_end = time.perf_counter_ns() + print(f"[INFO]: Scene creation time: {(setup_time_end - setup_time_begin) / 1e6:.2f} ms") + # Reset sim + reset_time_begin = time.perf_counter_ns() + sim.reset() + reset_time_end = time.perf_counter_ns() + print(f"[INFO]: Sim start time: {(reset_time_end - reset_time_begin) / 1e6:.2f} ms") + return sim, scene, sim.get_physics_dt() + + +def _run_benchmark(scene_cfg: RayCasterBenchmarkSceneCfg, sensor_name: str): + sim, scene, sim_dt = _setup_scene(scene_cfg) + sensor: RayCaster = scene[sensor_name] + # Warmup + for _ in range(args_cli.warmup): + sim.step() + sensor.update(dt=sim_dt, force_recompute=True) + + used_memory = 0.0 + + # Timing + t0 = time.perf_counter_ns() + for _ in range(args_cli.steps): + sim.step() + sensor.update(dt=sim_dt, force_recompute=True) + free, total = torch.cuda.mem_get_info(args_cli.device) + used_memory += (total - free) / 1024**2 # Convert to MB + t1 = time.perf_counter_ns() + per_step_ms = (t1 - t0) / args_cli.steps / 1e6 + avg_memory = used_memory / args_cli.steps + # Cleanup + # stop simulation + # note: cannot use self.sim.stop() since it does one render step after stopping!! This doesn't make sense :( + sim._timeline.stop() + # clear the stage + sim.clear_all_callbacks() + sim.clear_instance() + SimulationManager._simulation_manager_interface.reset() + SimulationManager._callbacks.clear() + + return { + "rays_per_sensor": int(sensor.num_rays), + "total_rays": int(sensor.num_rays * sensor.num_instances), + "per_step_ms": float(per_step_ms), + "avg_memory": float(avg_memory), + "num_meshes": len(np.unique([m.id for m in sensor.meshes.values()])), + } + + +def main(): + """Main function.""" + # Prepare benchmark + + # BENCHMARK 1 - Compare Single VS Multi + + print("=== Benchmarking Multi vs Single Raycaster ===") + results: list[dict[str, object]] = [] + device_name = ( + torch.cuda.get_device_name(torch.cuda.current_device()) if torch.cuda.is_available() else platform.processor() + ) + + _MESH_CONVERTERS_CALLBACKS["Sphere"] = lambda p: _create_sphere_trimesh(p, subdivisions=5) + # Compare multi mesh performance over different number of assets. + # More specifically, compare reference vs non-reference meshes and their respective memory usage. + for idx, num_assets in enumerate(NUM_ASSETS_MEMORY): + for reference_meshes in [True]: + if num_assets > 16 and not reference_meshes: + continue # Skip this, otherwise we run out of memory + + print(f"\n[INFO]: Benchmarking with {num_assets} assets. {idx} / {len(NUM_ASSETS_MEMORY)}") + num_envs = 1024 + resolution = 0.1 + multi_scene_cfg = _make_scene_cfg_multi( + num_envs=num_envs, + resolution=resolution, + debug_vis=not args_cli.headless, + track_mesh_transforms=True, + num_assets=num_assets, + reference_meshes=reference_meshes, + ) + result = _run_benchmark(multi_scene_cfg, "height_scanner_multi") + result["num_envs"] = num_envs + result["resolution"] = resolution + result["mode"] = "multi" + result["reference_meshes"] = reference_meshes + result["num_assets"] = num_assets + + print(result) + results.append(result) + del multi_scene_cfg + + df_num_assets = pd.DataFrame(results) + df_num_assets["device"] = device_name + os.makedirs(OUTPUT_DIR, exist_ok=True) + df_num_assets.to_csv(os.path.join(OUTPUT_DIR, "ray_caster_benchmark_num_assets_reference.csv"), index=False) + + results: list[dict[str, object]] = [] + + for idx, num_envs in enumerate(NUM_ENVS): + print(f"\n[INFO]: Benchmarking with {num_envs} envs. {idx + 1} / {len(NUM_ENVS)}") + + # Default Raycaster + for resolution in RESOLUTIONS: + single_scene_cfg = _make_scene_cfg_single( + num_envs=num_envs, + resolution=resolution, + debug_vis=not args_cli.headless, + ) + result = _run_benchmark(single_scene_cfg, "height_scanner") + result["num_envs"] = num_envs + result["resolution"] = resolution + result["mode"] = "single" + result["num_assets"] = 0 + results.append(result) + del single_scene_cfg + + # Multi Raycaster + for resolution in RESOLUTIONS: + multi_scene_cfg = _make_scene_cfg_multi( + num_envs=num_envs, + resolution=resolution, + debug_vis=not args_cli.headless, + track_mesh_transforms=False, + num_assets=0, + ) + result = _run_benchmark(multi_scene_cfg, "height_scanner_multi") + result["num_envs"] = num_envs + result["resolution"] = resolution + result["mode"] = "multi" + result["num_assets"] = 0 + results.append(result) + del multi_scene_cfg + + df_single_vs_multi = pd.DataFrame(results) + df_single_vs_multi["device"] = device_name + os.makedirs(OUTPUT_DIR, exist_ok=True) + df_single_vs_multi.to_csv(os.path.join(OUTPUT_DIR, "ray_caster_benchmark_single_vs_multi.csv"), index=False) + + print("\n=== Benchmarking Multi Raycaster with different number of assets and faces ===") + results: list[dict[str, object]] = [] + + # Keep fixed resolution for the subdivision and num assets benchmarks + resolution = 0.05 + + # Compare multi mesh performance over different number of assets + for idx, num_assets in enumerate(NUM_ASSETS): + print(f"\n[INFO]: Benchmarking with {num_assets} assets. {idx} / {len(NUM_ASSETS)}") + multi_scene_cfg = _make_scene_cfg_multi( + num_envs=num_envs, + resolution=resolution, + debug_vis=not args_cli.headless, + track_mesh_transforms=True, + num_assets=num_assets, + ) + result = _run_benchmark(multi_scene_cfg, "height_scanner_multi") + result["num_envs"] = num_envs + result["resolution"] = resolution + result["mode"] = "multi" + result["num_assets"] = num_assets + results.append(result) + del multi_scene_cfg + + df_num_assets = pd.DataFrame(results) + df_num_assets["device"] = device_name + df_num_assets.to_csv(os.path.join(OUTPUT_DIR, "ray_caster_benchmark_num_assets.csv"), index=False) + + print("\n=== Benchmarking Multi Raycaster with different number of faces ===") + results: list[dict[str, object]] = [] + # Compare multi mesh performance over different number of vertices + for idx, subdivision in enumerate(MESH_SUBDIVISIONS): + print(f"\n[INFO]: Benchmarking with {subdivision} subdivisions. {idx} / {len(MESH_SUBDIVISIONS)}") + _MESH_CONVERTERS_CALLBACKS["Sphere"] = lambda p: _create_sphere_trimesh(p, subdivisions=subdivision) + multi_scene_cfg = _make_scene_cfg_multi( + num_envs=num_envs, + resolution=resolution, + debug_vis=not args_cli.headless, + track_mesh_transforms=False, # Only static ground + num_assets=1, + ) + result = _run_benchmark(multi_scene_cfg, "height_scanner_multi") + result["num_envs"] = num_envs + result["resolution"] = resolution + result["mode"] = "multi" + result["num_assets"] = 1 + result["num_faces"] = 20 * (4**subdivision) + results.append(result) + del multi_scene_cfg + + df_num_faces = pd.DataFrame(results) + df_num_faces["device"] = device_name + df_num_faces.to_csv(os.path.join(OUTPUT_DIR, "ray_caster_benchmark_num_faces.csv"), index=False) + + # Create .md file with all three tables + for df, title in zip( + [df_single_vs_multi, df_num_assets, df_num_faces], ["Single vs Multi", "Num Assets", "Num Faces"] + ): + with open(os.path.join(OUTPUT_DIR, f"ray_caster_benchmark_{title}.md"), "w") as f: + f.write(f"# {title}\n\n") + f.write(dataframe_to_markdown(df, floatfmt=".3f")) + f.write("\n\n") + + +if __name__ == "__main__": + # run the main function + main() + # close sim app + simulation_app.close() diff --git a/scripts/benchmarks/benchmark_ray_caster_cli.py b/scripts/benchmarks/benchmark_ray_caster_cli.py new file mode 100644 index 00000000000..52912f4ed17 --- /dev/null +++ b/scripts/benchmarks/benchmark_ray_caster_cli.py @@ -0,0 +1,396 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Benchmark ray-caster sensors (single-mesh and multi-mesh). + +This script creates a simple scene with: +- a ground plane +- 0 to N moving spheres as rigid bodies + +It then runs three different benchmarks: +1. Single vs Multi mesh ray caster against ground only +2. Multi-mesh ray caster against ground + various number of spheres +3. Multi-mesh ray caster against spheres with different numbers of vertices +""" + +"""Launch Isaac Sim Simulator first.""" + +import argparse +import numpy as np +import os +import platform +import time +import torch + +import pandas as pd + +from isaaclab.app import AppLauncher + +# add argparse arguments +parser = argparse.ArgumentParser(description="Benchmark ray caster sensors.") +parser.add_argument("--steps", type=int, default=1000, help="Steps per resolution for timing.") +parser.add_argument("--warmup", type=int, default=50, help="Warmup steps before timing.") +parser.add_argument("--task", type=str, default="ray_caster_benchmark_single_vs_multi", help="Name of the task.") +parser.add_argument("--dirname", type=str, default="default", help="Output directory name (overrides timestamp).") +parser.add_argument("--num_assets", type=int, default=1, help="Number of assets (spheres) per env for multi-mesh benchmarks.") +parser.add_argument("--num_envs", type=int, default=1024, help="Number of envs to simulate.") +parser.add_argument("--mesh_subdivisions", type=int, default=2, help="Number of subdivisions for sphere mesh complexity.") +parser.add_argument("--resolution", type=float, default=0.05, help="Ray caster resolution for multi-mesh benchmarks.") +parser.add_argument("--raycaster_type", type=str, default="multi", help="Type of ray caster to benchmark: single or multi.", choices=["single", "multi"]) + +# TODO MOVE OUTSIDE +# # Num assets for benchmarking memory usage with and without caching. +# NUM_ASSETS_MEMORY = [1, 2, 4, 8, 16, 32] +# # Num assets for benchmarking scaling performance of multi-mesh ray caster. +# NUM_ASSETS = [0, 1, 2, 4, 8, 16, 32] +# # Num envs for benchmarking single vs multi mesh ray caster. +# NUM_ENVS = [32, 64, 128, 256, 512, 1024, 2048, 4096] +# # Num subdivisions for benchmarking mesh complexity. +# MESH_SUBDIVISIONS = [0, 1, 2, 3, 4, 5] +# # Different ray caster resolutions to benchmark. Num rays will be (5 / res)^2, e.g. 625, 2500, 10000, 11111 +# RESOLUTIONS: list[float] = [0.2, 0.1, 0.05, 0.015] + + + +# append AppLauncher cli args +AppLauncher.add_app_launcher_args(parser) +# parse the arguments +args_cli, _ = parser.parse_known_args() +args_cli.headless = True + +# launch omniverse app +app_launcher = AppLauncher(args_cli) +simulation_app = app_launcher.app + +"""Rest everything follows.""" + +import isaacsim.core.utils.stage as stage_utils +from isaacsim.core.simulation_manager import SimulationManager + +import isaaclab.sim as sim_utils +from isaaclab.assets import AssetBaseCfg, RigidObjectCfg +from isaaclab.scene import InteractiveScene, InteractiveSceneCfg +from isaaclab.sensors.ray_caster import MultiMeshRayCasterCfg, RayCaster, RayCasterCfg, patterns +from isaaclab.sim import SimulationContext +from isaaclab.utils import configclass +from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR +from isaaclab.utils.mesh import _MESH_CONVERTERS_CALLBACKS, _create_sphere_trimesh + + +OUTPUT_DIR = "outputs/benchmarks/" + args_cli.dirname + + +@configclass +class RayCasterBenchmarkSceneCfg(InteractiveSceneCfg): + """Scene config with ground, robot, and optional spheres per env.""" + + # ground plane (rough) + ground = AssetBaseCfg( + prim_path="/World/ground", + spawn=sim_utils.UsdFileCfg(usd_path=f"{ISAAC_NUCLEUS_DIR}/Environments/Terrains/rough_plane.usd"), + ) + + # lights + dome_light = AssetBaseCfg( + prim_path="/World/Light", spawn=sim_utils.DomeLightCfg(intensity=3000.0, color=(0.75, 0.75, 0.75)) + ) + + robot: RigidObjectCfg = RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/ray_caster_origin", + spawn=sim_utils.SphereCfg( + radius=0.1, + rigid_props=sim_utils.RigidBodyPropertiesCfg(kinematic_enabled=True), + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.0, 0.6, 0.0)), + ), + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, 4.0)), + ) + + # spheres collection (optionally set at runtime) + spheres: RigidObjectCfg = RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/sphere", + spawn=sim_utils.SphereCfg( + radius=0.2, + rigid_props=sim_utils.RigidBodyPropertiesCfg(max_depenetration_velocity=1.0, disable_gravity=True), + mass_props=sim_utils.MassPropertiesCfg(mass=100.0), + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.5, 0.0, 0.0)), + collision_props=sim_utils.CollisionPropertiesCfg(), + ), + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, 1.0)), + ) + + # sensors (set at runtime) + height_scanner: RayCasterCfg | None = None + height_scanner_multi: MultiMeshRayCasterCfg | None = None + + +def _make_scene_cfg_single(num_envs: int, resolution: float, debug_vis: bool) -> RayCasterBenchmarkSceneCfg: + scene_cfg = RayCasterBenchmarkSceneCfg(num_envs=num_envs, env_spacing=2.0) + scene_cfg.height_scanner = RayCasterCfg( + prim_path="{ENV_REGEX_NS}/ray_caster_origin", + mesh_prim_paths=["/World/ground"], + pattern_cfg=patterns.GridPatternCfg(resolution=resolution, size=(5.0, 5.0)), + ray_alignment="world", + offset=RayCasterCfg.OffsetCfg(pos=(0.0, 0.0, 20.0)), + debug_vis=debug_vis, + ) + return scene_cfg + + +def _make_scene_cfg_multi( + num_envs: int, + resolution: float, + debug_vis: bool, + track_mesh_transforms: bool, + num_assets: int = 1, + reference_meshes: bool = True, +) -> RayCasterBenchmarkSceneCfg: + scene_cfg = RayCasterBenchmarkSceneCfg(num_envs=num_envs, env_spacing=2.0) + + obj_cfg = scene_cfg.spheres + if track_mesh_transforms: + # Enable gravity + obj_cfg.spawn.rigid_props.disable_gravity = False + + for i in range(num_assets): + new_obj_cfg = obj_cfg.replace(prim_path=f"{{ENV_REGEX_NS}}/sphere_{i}") + ratio = i / num_assets + new_obj_cfg.init_state.pos = (ratio - 0.5, ratio - 0.5, 1.0) + setattr(scene_cfg, f"sphere_{i}", new_obj_cfg) + del scene_cfg.spheres + + scene_cfg.height_scanner_multi = MultiMeshRayCasterCfg( + prim_path="{ENV_REGEX_NS}/ray_caster_origin", + mesh_prim_paths=["/World/ground"] + + [ + MultiMeshRayCasterCfg.RaycastTargetCfg( + target_prim_expr=f"/World/envs/env_.*/sphere_{i}", track_mesh_transforms=track_mesh_transforms + ) + for i in range(num_assets) + ], + pattern_cfg=patterns.GridPatternCfg(resolution=resolution, size=(5.0, 5.0)), + offset=RayCasterCfg.OffsetCfg(pos=(0.0, 0.0, 20.0)), + debug_vis=debug_vis, + ray_alignment="world", + reference_meshes=reference_meshes, + ) + return scene_cfg + + +def _setup_scene(scene_cfg: RayCasterBenchmarkSceneCfg) -> tuple[SimulationContext, InteractiveScene, float]: + # Create a new stage + stage_utils.create_new_stage() + # New simulation per run + sim_cfg = sim_utils.SimulationCfg(device=args_cli.device) + sim = SimulationContext(sim_cfg) + sim.set_camera_view((2.5, 0.0, 4.0), (0.0, 0.0, 2.0)) + # Create scene with sensors + setup_time_begin = time.perf_counter_ns() + scene = InteractiveScene(scene_cfg) + setup_time_end = time.perf_counter_ns() + print(f"[INFO]: Scene creation time: {(setup_time_end - setup_time_begin) / 1e6:.2f} ms") + # Reset sim + reset_time_begin = time.perf_counter_ns() + sim.reset() + reset_time_end = time.perf_counter_ns() + print(f"[INFO]: Sim start time: {(reset_time_end - reset_time_begin) / 1e6:.2f} ms") + return sim, scene, sim.get_physics_dt() + + +def _run_benchmark(scene_cfg: RayCasterBenchmarkSceneCfg, sensor_name: str): + sim, scene, sim_dt = _setup_scene(scene_cfg) + sensor: RayCaster = scene[sensor_name] + # Warmup + for _ in range(args_cli.warmup): + sim.step() + sensor.update(dt=sim_dt, force_recompute=True) + + used_memory = 0.0 + + # Timing + t0 = time.perf_counter_ns() + for _ in range(args_cli.steps): + sim.step() + sensor.update(dt=sim_dt, force_recompute=True) + free, total = torch.cuda.mem_get_info(args_cli.device) + used_memory += (total - free) / 1024**2 # Convert to MB + t1 = time.perf_counter_ns() + per_step_ms = (t1 - t0) / args_cli.steps / 1e6 + avg_memory = used_memory / args_cli.steps + # Cleanup + # stop simulation + # note: cannot use self.sim.stop() since it does one render step after stopping!! This doesn't make sense :( + sim._timeline.stop() + # clear the stage + sim.clear_all_callbacks() + sim.clear_instance() + SimulationManager._simulation_manager_interface.reset() + SimulationManager._callbacks.clear() + + return { + "rays_per_sensor": int(sensor.num_rays), + "total_rays": int(sensor.num_rays * sensor.num_instances), + "per_step_ms": float(per_step_ms), + "avg_memory": float(avg_memory), + "num_meshes": len(np.unique([m.id for m in sensor.meshes.values()])), + } + + +def main(): + """Main function.""" + # Prepare benchmark + + # BENCHMARK 1 - Compare Single VS Multi + + num_envs = args_cli.num_envs + resolution = args_cli.resolution + num_assets = args_cli.num_assets + mesh_subdivisions = args_cli.mesh_subdivisions + raycaster_type = args_cli.raycaster_type + print(f"[INFO]: Using {num_envs} envs, {num_assets} assets, {mesh_subdivisions} subdivisions, {resolution} resolution") + + device_name = ( + torch.cuda.get_device_name(torch.cuda.current_device()) if torch.cuda.is_available() else platform.processor() + ) + + if args_cli.task == "ray_caster_benchmark_num_assets_reference": + print("=== Benchmarking Num Assets ===") + results: list[dict[str, object]] = [] + + _MESH_CONVERTERS_CALLBACKS["Sphere"] = lambda p: _create_sphere_trimesh(p, subdivisions=5) + # Compare multi mesh performance over different number of assets. + # More specifically, compare reference vs non-reference meshes and their respective memory usage. + num_envs = 1024 + resolution = 0.1 + multi_scene_cfg = _make_scene_cfg_multi( + num_envs=num_envs, + resolution=resolution, + debug_vis=not args_cli.headless, + track_mesh_transforms=True, + num_assets=num_assets, + reference_meshes=True, + ) + result = _run_benchmark(multi_scene_cfg, "height_scanner_multi") + result["num_envs"] = num_envs + result["resolution"] = resolution + result["mode"] = "multi" + result["reference_meshes"] = True + result["num_assets"] = num_assets + + print(result) + results.append(result) + del multi_scene_cfg + + df_num_assets = pd.DataFrame(results) + df_num_assets["device"] = device_name + os.makedirs(OUTPUT_DIR, exist_ok=True) + if os.path.exists(os.path.join(OUTPUT_DIR, "ray_caster_benchmark_num_assets_reference.csv")): + df_existing = pd.read_csv(os.path.join(OUTPUT_DIR, "ray_caster_benchmark_num_assets_reference.csv")) + df_num_assets = pd.concat([df_existing, df_num_assets], ignore_index=True) + df_num_assets.to_csv(os.path.join(OUTPUT_DIR, "ray_caster_benchmark_num_assets_reference.csv"), index=False) + + if args_cli.task == "ray_caster_benchmark_single_vs_multi": + results: list[dict[str, object]] = [] + if raycaster_type == "single": + # Default Raycaster + single_scene_cfg = _make_scene_cfg_single( + num_envs=num_envs, + resolution=resolution, + debug_vis=not args_cli.headless, + ) + result = _run_benchmark(single_scene_cfg, "height_scanner") + result["num_envs"] = num_envs + result["resolution"] = resolution + result["mode"] = "single" + result["num_assets"] = 0 + results.append(result) + del single_scene_cfg + + if raycaster_type == "multi": + # Multi Raycaster + multi_scene_cfg = _make_scene_cfg_multi( + num_envs=num_envs, + resolution=resolution, + debug_vis=not args_cli.headless, + track_mesh_transforms=False, + num_assets=0, + ) + result = _run_benchmark(multi_scene_cfg, "height_scanner_multi") + result["num_envs"] = num_envs + result["resolution"] = resolution + result["mode"] = "multi" + result["num_assets"] = 0 + results.append(result) + del multi_scene_cfg + + df_single_vs_multi = pd.DataFrame(results) + df_single_vs_multi["device"] = device_name + os.makedirs(OUTPUT_DIR, exist_ok=True) + if os.path.exists(os.path.join(OUTPUT_DIR, "ray_caster_benchmark_single_vs_multi.csv")): + df_existing = pd.read_csv(os.path.join(OUTPUT_DIR, "ray_caster_benchmark_single_vs_multi.csv")) + df_single_vs_multi = pd.concat([df_existing, df_single_vs_multi], ignore_index=True) + df_single_vs_multi.to_csv(os.path.join(OUTPUT_DIR, "ray_caster_benchmark_single_vs_multi.csv"), index=False) + + if args_cli.task == "ray_caster_benchmark_num_assets_and_faces": + print("\n=== Benchmarking Multi Raycaster with different number of assets and faces ===") + results: list[dict[str, object]] = [] + # Compare multi mesh performance over different number of assets + multi_scene_cfg = _make_scene_cfg_multi( + num_envs=num_envs, + resolution=resolution, + debug_vis=not args_cli.headless, + track_mesh_transforms=True, + num_assets=num_assets, + ) + result = _run_benchmark(multi_scene_cfg, "height_scanner_multi") + result["num_envs"] = num_envs + result["resolution"] = resolution + result["mode"] = "multi" + result["num_assets"] = num_assets + results.append(result) + del multi_scene_cfg + + df_num_assets = pd.DataFrame(results) + df_num_assets["device"] = device_name + # check if df frame exists + if os.path.exists(os.path.join(OUTPUT_DIR, "ray_caster_benchmark_num_assets.csv")): + df_existing = pd.read_csv(os.path.join(OUTPUT_DIR, "ray_caster_benchmark_num_assets.csv")) + df_num_assets = pd.concat([df_existing, df_num_assets], ignore_index=True) + df_num_assets.to_csv(os.path.join(OUTPUT_DIR, "ray_caster_benchmark_num_assets.csv"), index=False) + + if args_cli.task == "ray_caster_benchmark_num_faces": + print("\n=== Benchmarking Multi Raycaster with different number of faces ===") + results: list[dict[str, object]] = [] + # Compare multi mesh performance over different number of vertices + _MESH_CONVERTERS_CALLBACKS["Sphere"] = lambda p: _create_sphere_trimesh(p, subdivisions=mesh_subdivisions) + multi_scene_cfg = _make_scene_cfg_multi( + num_envs=num_envs, + resolution=resolution, + debug_vis=not args_cli.headless, + track_mesh_transforms=False, # Only static ground + num_assets=1, + ) + result = _run_benchmark(multi_scene_cfg, "height_scanner_multi") + result["num_envs"] = num_envs + result["resolution"] = resolution + result["mode"] = "multi" + result["num_assets"] = 1 + result["num_faces"] = 20 * (4**mesh_subdivisions) + results.append(result) + del multi_scene_cfg + + df_num_faces = pd.DataFrame(results) + df_num_faces["device"] = device_name + # check if df frame exists + if os.path.exists(os.path.join(OUTPUT_DIR, "ray_caster_benchmark_num_faces.csv")): + df_existing = pd.read_csv(os.path.join(OUTPUT_DIR, "ray_caster_benchmark_num_faces.csv")) + df_num_faces = pd.concat([df_existing, df_num_faces], ignore_index=True) + + df_num_faces.to_csv(os.path.join(OUTPUT_DIR, "ray_caster_benchmark_num_faces.csv"), index=False) + +if __name__ == "__main__": + # run the main function + main() + # close sim app + simulation_app.close() diff --git a/scripts/benchmarks/compare_raycasters.sh b/scripts/benchmarks/compare_raycasters.sh new file mode 100755 index 00000000000..9db9907b1a7 --- /dev/null +++ b/scripts/benchmarks/compare_raycasters.sh @@ -0,0 +1,8 @@ +#!/bin/bash + +# Expects first argument api key, second wandb entity (optional) +export WANDB_API_KEY=$1 +export WANDB_ENTITY=$2 +echo Running raycaster benchmarks with wandb entity $WANDB_ENTITY +${ISAACLAB_PATH}/_isaac_sim/python.sh ${ISAACLAB_PATH}/scripts/benchmarks/benchmark_ray_caster.py +${ISAACLAB_PATH}/_isaac_sim/python.sh ${ISAACLAB_PATH}/scripts/benchmarks/sync_raycast_results.py diff --git a/scripts/benchmarks/local_utils.py b/scripts/benchmarks/local_utils.py new file mode 100644 index 00000000000..0b66508b452 --- /dev/null +++ b/scripts/benchmarks/local_utils.py @@ -0,0 +1,60 @@ + +import math +import pandas as pd + +################### +# Output Printing # +################### + +def dataframe_to_markdown( + df: pd.DataFrame, index: bool = True, floatfmt: str | None = ".3f", nan_rep: str = "", align: str = "left" +) -> str: + """ + Convert a pandas DataFrame to a Markdown table (no extra deps). + - index: include the index as the first column + - floatfmt: e.g. '.2f' for floats; set None to use str() directly + - nan_rep: string to show for NaNs/None + - align: 'left' | 'center' | 'right' + """ + + def _fmt(x): + if x is None or (isinstance(x, float) and math.isnan(x)): + return nan_rep + if isinstance(x, float) and floatfmt is not None: + return format(x, floatfmt) + return str(x) + + def _esc(s: str) -> str: + # Escape pipes so they don't break the Markdown table + return s.replace("|", r"\|") + + # Build header and rows + headers = [str(c) for c in df.columns] + rows = [[_fmt(v) for v in row] for row in df.to_numpy().tolist()] + + if index: + headers = [df.index.name or ""] + headers + idx_col = [str(i) for i in df.index.tolist()] + rows = [[idx] + r for idx, r in zip(idx_col, rows)] + + # Compute column widths + cols = list(zip(*([headers] + rows))) if headers else [] + widths = [max(len(_esc(h)), *(len(_esc(cell)) for cell in col)) for h, col in zip(headers, cols)] + + # Alignment rule + def rule(w): + if align == "right": + return "-" * (w - 1) + ":" + if align == "center": + return ":" + "-" * (w - 2 if w > 2 else 1) + ":" + return ":" + "-" * (w - 1) # left + + # Build markdown lines + def fmt_row(cells): + return "| " + " | ".join(_esc(c).ljust(w) for c, w in zip(cells, widths)) + " |" + + header_line = fmt_row(headers) + sep_line = "| " + " | ".join(rule(w) for w in widths) + " |" + body_lines = [fmt_row(r) for r in rows] + + return "\n".join([header_line, sep_line, *body_lines]) \ No newline at end of file diff --git a/scripts/benchmarks/plot_environments.py b/scripts/benchmarks/plot_environments.py new file mode 100644 index 00000000000..857c9b60a10 --- /dev/null +++ b/scripts/benchmarks/plot_environments.py @@ -0,0 +1,154 @@ +import pandas as pd +import matplotlib.pyplot as plt +import scienceplots + +plt.style.use(['science']) +SMALL_SIZE = 24 +MEDIUM_SIZE = 28 +LEGEND_SIZE = 24 +BIGGER_SIZE = 32 + +plt.rc('font', size=20) # controls default text sizes +plt.rc('axes', titlesize=SMALL_SIZE) # fontsize of the axes title +plt.rc('axes', labelsize=MEDIUM_SIZE) # fontsize of the x and y labels +plt.rc('xtick', labelsize=SMALL_SIZE) # fontsize of the tick labels +plt.rc('ytick', labelsize=SMALL_SIZE) # fontsize of the tick labels +plt.rc('legend', fontsize=LEGEND_SIZE) # legend fontsize +plt.rc('figure', titlesize=BIGGER_SIZE) # fontsize of the figure title +plt.rc("lines", linewidth=2) + +MACHINE_NBR_GPUS = "RTX PRO 6000" + +NBR_GPUS_COLORS = { # RTX PRO 6000 + "Single-GPU": "#d08770", + "2-GPU": "#ebcb8b", + "4-GPU": "#a3be8c", + "8-GPU": "#b48ead" +} + +MACHINE_NBR_GPUS_COLORS = { # L40 + 5090 + "L40": "#5e81ac", + "5090": "#a3be8c" +} + + +# 1) Load CSV (no headers so we can custom-parse) +file_path = "/home/pascalr/Downloads/Isaac Lab Whitepaper Benchmarks - benchmark result.csv" +raw = pd.read_csv(file_path, header=None) + +# 2) Locate the "Mean" and "Std" block boundaries in row 0 +row0 = raw.iloc[0].astype(str) +mean_start = row0[row0.str.contains("Mean", na=False)].index[0] +std_start = row0[row0.str.contains("Std", na=False)].index[0] +end_col = raw.shape[1] + +# Helper to parse a block (Mean or Std) into tidy format +def parse_block(start_col, end_col, value_name): + gpu_headers = raw.iloc[1] # GPU names line (e.g., "Single-GPU", "2-GPU") + env_headers = raw.iloc[2] # Environment counts line (1024..16384) + data = raw.iloc[3:].reset_index(drop=True) + data = data.rename(columns={0: "Machine", 1: "Task"}) + + current_gpu = None + records = [] + for col in range(start_col, end_col): + gpu = gpu_headers[col] + env = env_headers[col] + + # Update GPU group when encountered + if isinstance(gpu, str) and "GPU" in gpu: + current_gpu = gpu + + if current_gpu is None: + continue + + # Parse environment count + try: + env_val = int(float(env)) + except Exception: + continue + + # Collect values per row (task/machine) + for i in range(len(data)): + v = data.iloc[i, col] + if pd.isna(v): + continue + records.append({ + "Machine": data.loc[i, "Machine"], + "Task": data.loc[i, "Task"], + "GPU_Config": current_gpu, + "Envs": env_val, + value_name: float(v) + }) + return pd.DataFrame(records) + +# 3) Parse mean and std blocks +mean_df = parse_block(mean_start, std_start, "Mean") +std_df = parse_block(std_start, end_col, "Std") + +# 4) Merge into one tidy dataframe +tidy = pd.merge( + mean_df, + std_df, + on=["Machine", "Task", "GPU_Config", "Envs"], + how="left" +) + +# 5) Ensure numeric types +tidy["Envs"] = pd.to_numeric(tidy["Envs"], errors="coerce").astype("Int64") +tidy["Mean"] = pd.to_numeric(tidy["Mean"], errors="coerce") +tidy["Std"] = pd.to_numeric(tidy["Std"], errors="coerce") +tidy = tidy.dropna(subset=["Envs", "Mean"]) + +# 6) Plot settings +tasks = tidy["Task"].dropna().unique() +machines = tidy["Machine"].dropna().unique() + +for i, task in enumerate(tasks): + plt.figure(figsize=(12, 8), dpi=300) + sub = tidy[tidy["Task"] == task] + + x_ticks = set() + + # --- Plot L40 + L40S (Single-GPU only) --- + for machine, color in MACHINE_NBR_GPUS_COLORS.items(): + g = sub[(sub["Machine"] == machine) & (sub["GPU_Config"] == "Single-GPU")].sort_values("Envs") + if not g.empty: + x = g["Envs"].astype(int) + y = g["Mean"] + yerr = g["Std"].fillna(0) + plt.plot(x, y, label=f"{machine} - Single-GPU", linestyle="dashdot", marker="o", color=color) + plt.fill_between(x, y - yerr, y + yerr, color=color, alpha=0.2) + + x_ticks.update(x) + + # --- Plot RTX PRO 6000 (all GPU configs) --- + for gpu, color in NBR_GPUS_COLORS.items(): + g = sub[(sub["Machine"] == "RTX PRO 6000") & (sub["GPU_Config"] == gpu)].sort_values("Envs") + if not g.empty: + x = g["Envs"].astype(int) + y = g["Mean"] + yerr = g["Std"].fillna(0) + plt.plot(x, y, label=f"RTX PRO 6000 - {gpu}", color=NBR_GPUS_COLORS[gpu], marker="o") + plt.fill_between(x, y - yerr, y + yerr, color=NBR_GPUS_COLORS[gpu], alpha=0.2) + + x_ticks.update(x) + + # Axes styling + plt.xscale("log", base=2) + plt.yscale("log") + plt.xticks(list(x_ticks), [str(x) for x in x_ticks]) + plt.xlabel("Number of Environments") + plt.legend( + loc="upper left", + frameon=True, # draw a box around the legend + facecolor='white', # white background + edgecolor='black', # optional: black border + framealpha=0.8 # transparency of the box (0.0 to 1.0) + ) + plt.ylabel("FPS") + # plt.title(task, weight="bold") + + plt.tight_layout(pad=2.0) + plt.savefig(f"benchmark_{task}_gpu_model_nbr_comparison.png", dpi=300) + plt.close() diff --git a/scripts/benchmarks/plot_perceptive_environment.py b/scripts/benchmarks/plot_perceptive_environment.py new file mode 100644 index 00000000000..fbe893243cd --- /dev/null +++ b/scripts/benchmarks/plot_perceptive_environment.py @@ -0,0 +1,192 @@ +import pandas as pd +import matplotlib.pyplot as plt +import scienceplots + +plt.style.use(['science']) +SMALL_SIZE = 24 +MEDIUM_SIZE = 28 +LEGEND_SIZE = 20 +BIGGER_SIZE = 32 + +plt.rc('font', size=20) # controls default text sizes +plt.rc('axes', titlesize=SMALL_SIZE) # fontsize of the axes title +plt.rc('axes', labelsize=MEDIUM_SIZE) # fontsize of the x and y labels +plt.rc('xtick', labelsize=SMALL_SIZE) # fontsize of the tick labels +plt.rc('ytick', labelsize=SMALL_SIZE) # fontsize of the tick labels +plt.rc('legend', fontsize=LEGEND_SIZE) # legend fontsize +plt.rc('figure', titlesize=BIGGER_SIZE) # fontsize of the figure title +plt.rc("lines", linewidth=2) + +MACHINE_NBR_GPUS = "RTX PRO 6000" + +NBR_GPUS_COLORS = { + "Single-GPU": "#d08770", + "2-GPU": "#ebcb8b", + "4-GPU": "#a3be8c", + "8-GPU": "#b48ead" +} + +MACHINE_NBR_GPUS_COLORS = { + "L40": "#5e81ac", + "5090": "#a3be8c", + "RTX PRO 6000": "#d08770" +} + +COMPARISON_ENVIRONMENTS = [ + {"tiled": "Isaac-Dexsuite-Kuka-Allegro-Lift-Depth-TiledCamera-v0", "raycaster": "Isaac-Dexsuite-Kuka-Allegro-Lift-Depth-RayCasterCamera-v0"}, + {"raycaster": "Isaac-Navigation-Flat-Anymal-C-v0"}, +] +COMPARISON_ENVIRONMENTS_NAME = ["dexsuite", "navigation"] + +# 1) Load CSV (no headers so we can custom-parse) +file_path = "/home/pascalr/Downloads/Isaac Lab Whitepaper Benchmarks - benchmark result.csv" +raw = pd.read_csv(file_path, header=None) + +# 2) Locate the "Mean" and "Std" block boundaries in row 0 +row0 = raw.iloc[0].astype(str) +mean_start = row0[row0.str.contains("Mean", na=False)].index[0] +std_start = row0[row0.str.contains("Std", na=False)].index[0] +end_col = raw.shape[1] + +# Helper to parse a block (Mean or Std) into tidy format +def parse_block(start_col, end_col, value_name): + gpu_headers = raw.iloc[1] # GPU names line (e.g., "Single-GPU", "2-GPU") + env_headers = raw.iloc[2] # Environment counts line (1024..16384) + data = raw.iloc[3:].reset_index(drop=True) + data = data.rename(columns={0: "Machine", 1: "Task"}) + + current_gpu = None + records = [] + for col in range(start_col, end_col): + gpu = gpu_headers[col] + env = env_headers[col] + + # Update GPU group when encountered + if isinstance(gpu, str) and "GPU" in gpu: + current_gpu = gpu + + if current_gpu is None: + continue + + # Parse environment count + try: + env_val = int(float(env)) + except Exception: + continue + + # Collect values per row (task/machine) + for i in range(len(data)): + v = data.iloc[i, col] + if pd.isna(v): + continue + records.append({ + "Machine": data.loc[i, "Machine"], + "Task": data.loc[i, "Task"], + "GPU_Config": current_gpu, + "Envs": env_val, + value_name: float(v) + }) + return pd.DataFrame(records) + +# 3) Parse mean and std blocks +mean_df = parse_block(mean_start, std_start, "Mean") +std_df = parse_block(std_start, end_col, "Std") + +# 4) Merge into one tidy dataframe +tidy = pd.merge( + mean_df, + std_df, + on=["Machine", "Task", "GPU_Config", "Envs"], + how="left" +) + +# 5) Ensure numeric types +tidy["Envs"] = pd.to_numeric(tidy["Envs"], errors="coerce").astype("Int64") +tidy["Mean"] = pd.to_numeric(tidy["Mean"], errors="coerce") +tidy["Std"] = pd.to_numeric(tidy["Std"], errors="coerce") +tidy = tidy.dropna(subset=["Envs", "Mean"]) + +# 6) Plot settings +tasks = tidy["Task"].dropna().unique() +machines = tidy["Machine"].dropna().unique() + +for i, comparison_environment in enumerate(COMPARISON_ENVIRONMENTS): + fig, axs = plt.subplots(1, 2, figsize=(16, 10), sharey=False) + + x_ticks_nbr_gpus = set() + x_ticks_machine = set() + + for perception_mode, style in [("tiled", "dashed"), ("raycaster", "dotted")]: + if perception_mode not in comparison_environment: + continue + + sub = tidy[tidy["Task"] == comparison_environment[perception_mode]] + + for machine, color in MACHINE_NBR_GPUS_COLORS.items(): + g = sub[(sub["Machine"] == machine) & (sub["GPU_Config"] == "Single-GPU")].sort_values("Envs") + if not g.empty: + x = g["Envs"].astype(int) + y = g["Mean"] + yerr = g["Std"].fillna(0) + axs[0].plot(x, y, label=f"{machine} - {perception_mode}", linestyle=style, marker="o", color=color) + axs[0].fill_between(x, y - yerr, y + yerr, color=color, alpha=0.2) + x_ticks_machine.update(x) + + for gpu, color in NBR_GPUS_COLORS.items(): + g = sub[(sub["Machine"] == "RTX PRO 6000") & (sub["GPU_Config"] == gpu)].sort_values("Envs") + if not g.empty: + x = g["Envs"].astype(int) + y = g["Mean"] + yerr = g["Std"].fillna(0) + axs[1].plot(x, y, label=f"{gpu} - {perception_mode}", linestyle=style, color=NBR_GPUS_COLORS[gpu], marker="o") + axs[1].fill_between(x, y - yerr, y + yerr, color=NBR_GPUS_COLORS[gpu], alpha=0.2) + x_ticks_nbr_gpus.update(x) + + # Axes styling + axs[0].set_xscale("log", base=2) + axs[0].set_yscale("log") + axs[0].set_xticks(list(x_ticks_machine), [str(x) for x in x_ticks_machine]) + axs[1].set_xticks(list(x_ticks_nbr_gpus), [str(x) for x in x_ticks_nbr_gpus]) + axs[0].set_xlabel("Number of Environments") + axs[1].set_xlabel("Number of Environments") + axs[0].set_ylabel("FPS") + + # ===== Separate legends for each subplot, below the figure ===== + # Collect handles/labels from each subplot + handles0, labels0 = axs[0].get_legend_handles_labels() + handles1, labels1 = axs[1].get_legend_handles_labels() + + # Legend for left subplot + leg0 = axs[0].legend( + handles0, labels0, + loc="upper center", + bbox_to_anchor=(-0.1, -0.25, 1.2, 0), # full width of subplot + mode="expand", # expand across width + ncol=2, + frameon=True, + facecolor="white", + edgecolor="black", + framealpha=0.8, + fontsize=LEGEND_SIZE # slightly smaller font + ) + + # Legend for right subplot + leg1 = axs[1].legend( + handles1, labels1, + loc="upper center", + bbox_to_anchor=(-0.1, -0.25, 1.2, 0), # full width of subplot + mode="expand", + ncol=2, + frameon=True, + facecolor="white", + edgecolor="black", + framealpha=0.8, + fontsize=LEGEND_SIZE + ) + + # Extra space below for legends + plt.tight_layout(pad=2.0, rect=[0, 0.05, 1, 1]) + + plt.savefig(f"benchmark_{COMPARISON_ENVIRONMENTS_NAME[i]}_perceptive_environment.pdf", dpi=300, + bbox_extra_artists=(leg0, leg1), bbox_inches="tight") + plt.close() diff --git a/scripts/benchmarks/plot_raycast_results.py b/scripts/benchmarks/plot_raycast_results.py new file mode 100644 index 00000000000..e0d8e8635ef --- /dev/null +++ b/scripts/benchmarks/plot_raycast_results.py @@ -0,0 +1,186 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import matplotlib.pyplot as plt + +import pandas as pd + + +def compare_single_vs_multi(): + df = pd.read_csv("outputs/benchmarks/ray_caster_benchmark_single_vs_multi.csv") + # %% Types & cleaning + df["resolution"] = df["resolution"].astype(float) + df["num_envs"] = df["num_envs"].astype(int) + df["avg_memory"] = df["avg_memory"].astype(float) + df["time_per_ray_us"] = df["per_step_ms"] * 1000.0 / df["total_rays"] # µs / ray + df["rays_per_env"] = df["total_rays"] / df["num_envs"] + df["fps"] = 1.0 / (df["per_step_ms"] * 1e-3) + df["rays_per_us"] = 1.0 / df["time_per_ray_us"] + + unique_res = sorted(df["resolution"].unique()) + color_palette = plt.cm.tab10.colors + color_map = {res: color_palette[i % len(color_palette)] for i, res in enumerate(unique_res)} + dash_map = {"single": "-", "multi": "--"} + + fig, axes = plt.subplots(1, 3, figsize=(20, 10)) + + # 1. Steps per s vs Num Envs + for res in unique_res: + for mode in df["mode"].unique(): + sub_df = df[(df["resolution"] == res) & (df["mode"] == mode)] + axes[0].plot( + sub_df["num_envs"], + sub_df["fps"], + dash_map.get(mode, "-"), + color=color_map[res], + marker="o", + label=f"{mode} - Resolution={res}, Rays/Env={sub_df['rays_per_env'].values[0]}", + ) + axes[0].set_xlabel("Number of Environments") + axes[0].set_ylabel("FPS (frames per second)") + axes[0].set_title("FPS vs Num Envs") + axes[0].legend() + axes[0].grid(True) + + # 2. Rays per µs vs Num Envs + for res in unique_res: + for mode in df["mode"].unique(): + sub_df = df[(df["resolution"] == res) & (df["mode"] == mode)] + axes[1].plot( + sub_df["num_envs"], + sub_df["rays_per_us"], + dash_map.get(mode, "-"), + color=color_map[res], + marker="o", + label=f"{mode} - Resolution={res}, Rays/Env={sub_df['rays_per_env'].values[0]}", + ) + axes[1].set_xlabel("Number of Environments") + axes[1].set_ylabel("Rays per µs") + axes[1].set_title("Rays per µs vs Num Envs") + axes[1].legend() + axes[1].grid(True) + + # 3. VRAM usage vs Number of Envs + for res in unique_res: + for mode in df["mode"].unique(): + sub_df = df[(df["resolution"] == res) & (df["mode"] == mode)] + axes[2].plot( + sub_df["num_envs"], + sub_df["avg_memory"], + dash_map.get(mode, "-"), + color=color_map[res], + marker="o", + label=f"{mode} - Resolution={res}, Rays/Env={sub_df['rays_per_env'].values[0]}", + ) + axes[2].set_xlabel("Number of Environments") + axes[2].set_ylabel("Average VRAM Usage (MB)") + axes[2].set_title("Average VRAM Usage (MB) vs Num Envs") + axes[2].legend() + axes[2].grid(True) + fig.suptitle("Raycast Benchmark Comparison", fontsize=16) + fig.tight_layout(rect=[0, 0.03, 1, 0.97]) + + # Save as PNG + plt.savefig("outputs/benchmarks/raycast_benchmark_comparison.png") + print("Saved plot to outputs/benchmarks/raycast_benchmark_comparison.png") + + +def compare_num_assets_vram_cache(): + """Plots Average steps per ms vs Num assets and Avg memory vs num assets""" + df = pd.read_csv("outputs/benchmarks/ray_caster_benchmark_num_assets_reference.csv") + df["num_assets"] = df["num_assets"].astype(int) + df["avg_memory"] = df["avg_memory"].astype(float) + df["fps"] = 1.0 / (df["per_step_ms"] * 1e-3) + df["reference_meshes"] = df["reference_meshes"].astype(bool) + + fig, axes = plt.subplots(1, 1, figsize=(10, 10)) + # 2. VRAM usage vs Number of Assets + for ref in [True, False]: + sub_df = df[df["reference_meshes"] == ref] + axes.plot(sub_df["num_assets"], sub_df["avg_memory"], marker="o", label=f"Reference Meshes: {ref}") + axes.set_xlabel("Number of Assets") + axes.set_ylabel("Average VRAM Usage (MB)") + axes.set_title("Average VRAM Usage (MB) vs Num Assets") + axes.legend() + axes.grid(True) + + fig.suptitle("Raycast Benchmark Comparison (Num Assets)", fontsize=16) + fig.tight_layout(rect=[0, 0.03, 1, 0.97]) + + # Save as PNG + plt.savefig("outputs/benchmarks/raycast_benchmark_comparison_num_assets_vram.png") + print("Saved plot to outputs/benchmarks/raycast_benchmark_comparison_num_assets_vram.png") + + +def compare_num_assets(): + """Plots Average steps per ms vs Num assets and Avg memory vs num assets""" + df = pd.read_csv("outputs/benchmarks/ray_caster_benchmark_num_assets.csv") + df["num_assets"] = df["num_assets"].astype(int) + df["avg_memory"] = df["avg_memory"].astype(float) + df["fps"] = 1.0 / (df["per_step_ms"] * 1e-3) + + fig, axes = plt.subplots(1, 2, figsize=(10, 10)) + # Plot FPS vs Num assets + axes[0].plot(df["num_assets"], df["fps"], marker="o") + axes[0].set_xlabel("Number of Assets") + axes[0].set_ylabel("FPS") + axes[0].set_title("FPS vs Num Assets") + axes[0].grid(True) + + # 2. VRAM usage vs Number of Assets + axes[1].plot(df["num_assets"], df["avg_memory"], marker="o") + axes[1].set_xlabel("Number of Assets") + axes[1].set_ylabel("Average VRAM Usage (MB)") + axes[1].set_title("Average VRAM Usage (MB) vs Num Assets") + axes[1].grid(True) + + fig.suptitle("Raycast Benchmark Comparison (Num Assets)", fontsize=16) + fig.tight_layout(rect=[0, 0.03, 1, 0.97]) + + # Save as PNG + plt.savefig("outputs/benchmarks/raycast_benchmark_comparison_num_assets.png") + print("Saved plot to outputs/benchmarks/raycast_benchmark_comparison_num_assets.png") + + +def compare_num_faces(): + """Plots Average steps per ms vs Num faces and Avg memory vs num faces""" + df = pd.read_csv("outputs/benchmarks/ray_caster_benchmark_num_faces.csv") + # %% Types & cleaning + df["num_faces"] = df["num_faces"].astype(int) + df["avg_memory"] = df["avg_memory"].astype(float) + + df["fps"] = 1.0 / (df["per_step_ms"] * 1e-3) + + fig, axes = plt.subplots(1, 2, figsize=(10, 10)) + + # 1. Steps per ms vs Num Faces + axes[0].plot(df["num_faces"], df["fps"], color="blue", marker="o", label="FPS") + axes[0].set_xlabel("Number of Faces") + axes[0].set_ylabel("FPS (frames per second)") + axes[0].set_title("FPS vs Num Faces") + axes[0].legend() + axes[0].grid(True) + + # 2. VRAM usage vs Number of Faces + axes[1].plot(df["num_faces"], df["avg_memory"], color="red", marker="o", label="Average VRAM Usage") + axes[1].set_xlabel("Number of Faces") + axes[1].set_ylabel("Average VRAM Usage (MB)") + axes[1].set_title("Average VRAM Usage (MB) vs Num Faces") + axes[1].legend() + axes[1].grid(True) + + fig.suptitle("Raycast Benchmark Comparison (Num Faces)", fontsize=16) + fig.tight_layout(rect=[0, 0.03, 1, 0.97]) + + # Save as PNG + plt.savefig("outputs/benchmarks/raycast_benchmark_comparison_num_faces.png") + print("Saved plot to outputs/benchmarks/raycast_benchmark_comparison_num_faces.png") + + +if __name__ == "__main__": + compare_single_vs_multi() + compare_num_assets_vram_cache() + compare_num_assets() + compare_num_faces() diff --git a/scripts/benchmarks/raycaster.sh b/scripts/benchmarks/raycaster.sh new file mode 100644 index 00000000000..2d955ad64ec --- /dev/null +++ b/scripts/benchmarks/raycaster.sh @@ -0,0 +1,73 @@ +#!/bin/bash +set -euo pipefail + +PYTHON=/workspace/isaaclab/_isaac_sim/python.sh +SCRIPT=scripts/benchmarks/benchmark_ray_caster_cli.py + +# =============================== +# Benchmark configurations +# =============================== + +NUM_ASSETS_MEMORY=(1 2 4 8 16 32) +NUM_ASSETS=(0 1 2 4 8 16 32) +NUM_ENVS=(32 64 128 256 512 1024 2048 4096) +MESH_SUBDIVISIONS=(0 1 2 3 4 5) +RESOLUTIONS=(0.2 0.1 0.05 0.015) + +# =============================== +# Run Benchmarks +# =============================== + +echo "=== Running Benchmarks ===" + +# Benchmark 1: Num Assets Reference +for num_assets in "${NUM_ASSETS_MEMORY[@]}"; do + echo "[RUN]: Num Assets Reference | num_assets=${num_assets}" + $PYTHON $SCRIPT \ + --task ray_caster_benchmark_num_assets_reference \ + --num_assets $num_assets \ + --num_envs 1024 \ + --resolution 0.05 +done + +# Benchmark 2: Single vs Multi (over envs & resolutions) +for num_envs in "${NUM_ENVS[@]}"; do + for res in "${RESOLUTIONS[@]}"; do + echo "[RUN]: Single Raycaster | num_envs=${num_envs}, res=${res}" + $PYTHON $SCRIPT \ + --task ray_caster_benchmark_single_vs_multi \ + --raycaster_type single \ + --num_envs $num_envs \ + --resolution $res + + echo "[RUN]: Multi Raycaster | num_envs=${num_envs}, res=${res}" + $PYTHON $SCRIPT \ + --task ray_caster_benchmark_single_vs_multi \ + --raycaster_type multi \ + --num_envs $num_envs \ + --resolution $res + done +done + +# Benchmark 3: Num Assets +for num_assets in "${NUM_ASSETS[@]}"; do + echo "[RUN]: Num Assets | num_assets=${num_assets}" + $PYTHON $SCRIPT \ + --task ray_caster_benchmark_num_assets_and_faces \ + --num_assets $num_assets \ + --resolution 0.05 \ + --num_envs 1024 +done + +# Benchmark 4: Mesh Subdivisions (Num Faces) +for subdiv in "${MESH_SUBDIVISIONS[@]}"; do + echo "[RUN]: Mesh Subdivisions | subdiv=${subdiv}" + $PYTHON $SCRIPT \ + --task ray_caster_benchmark_num_faces \ + --mesh_subdivisions $subdiv \ + --resolution 0.05 \ + --num_envs 1024 \ + --num_assets 1 +done + +echo "=== All Benchmarks Completed ===" diff --git a/scripts/benchmarks/sync_raycast_results.py b/scripts/benchmarks/sync_raycast_results.py new file mode 100644 index 00000000000..d1492b89305 --- /dev/null +++ b/scripts/benchmarks/sync_raycast_results.py @@ -0,0 +1,171 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import os + +import pandas as pd +import plotly.express as px + +import wandb + + +def compare_single_vs_multi(): + csv_path = "outputs/benchmarks/ray_caster_benchmark_single_vs_multi.csv" + df = pd.read_csv(csv_path) + + artifact = wandb.Artifact("ray_caster_benchmark_single_vs_multi", type="dataset") + artifact.add_file(csv_path) + wandb.log_artifact(artifact) + + df["resolution"] = df["resolution"].astype(float) + df["num_envs"] = df["num_envs"].astype(int) + df["avg_memory"] = df["avg_memory"].astype(float) + df["time_per_ray_us"] = df["per_step_ms"] * 1000.0 / df["total_rays"] + df["rays_per_env"] = df["total_rays"] / df["num_envs"] + df["fps"] = 1.0 / (df["per_step_ms"] * 1e-3) + df["rays_per_us"] = 1.0 / df["time_per_ray_us"] + + df["resolution_label"] = df.apply(lambda r: f"{r['mode']} - {(5.0 / r['resolution'])**2:.2f} rays", axis=1) + + fig_fps = px.line( + df, + x="num_envs", + y="fps", + color="resolution_label", + markers=True, + title="FPS vs Num Envs", + ) + + fig_rays = px.line( + df, + x="num_envs", + y="rays_per_us", + color="resolution_label", + markers=True, + title="Rays per µs vs Num Envs", + ) + + fig_mem = px.line( + df, + x="num_envs", + y="avg_memory", + color="resolution_label", + markers=True, + title="Average VRAM Usage (MB) vs Num Envs", + ) + + wandb.log({ + "FPS vs Num Envs": fig_fps, + "Rays per µs vs Num Envs": fig_rays, + "VRAM vs Num Envs": fig_mem, + "Single vs Multi Table": wandb.Table(dataframe=df), + }) + + +def compare_num_assets_vram_cache(): + csv_path = "outputs/benchmarks/ray_caster_benchmark_num_assets_reference.csv" + df = pd.read_csv(csv_path) + + artifact = wandb.Artifact("ray_caster_benchmark_num_assets_reference", type="dataset") + artifact.add_file(csv_path) + wandb.log_artifact(artifact) + + df["num_assets"] = df["num_assets"].astype(int) + df["avg_memory"] = df["avg_memory"].astype(float) + df["fps"] = 1.0 / (df["per_step_ms"] * 1e-3) + df["reference_meshes"] = df["reference_meshes"].astype(bool) + + fig = px.line( + df, + x="num_assets", + y="avg_memory", + color="reference_meshes", + markers=True, + title="Average VRAM Usage (MB) vs Num Assets", + ) + + wandb.log({ + "Num Assets (VRAM Cache)": fig, + "Num Assets VRAM Table": wandb.Table(dataframe=df), + }) + + +def compare_num_assets(): + csv_path = "outputs/benchmarks/ray_caster_benchmark_num_assets.csv" + df = pd.read_csv(csv_path) + + artifact = wandb.Artifact("ray_caster_benchmark_num_assets", type="dataset") + artifact.add_file(csv_path) + wandb.log_artifact(artifact) + + df["num_assets"] = df["num_assets"].astype(int) + df["avg_memory"] = df["avg_memory"].astype(float) + df["fps"] = 1.0 / (df["per_step_ms"] * 1e-3) + + fig_fps = px.line( + df, + x="num_assets", + y="fps", + markers=True, + title="FPS vs Num Assets", + ) + + fig_mem = px.line( + df, + x="num_assets", + y="avg_memory", + markers=True, + title="Average VRAM Usage (MB) vs Num Assets", + ) + + wandb.log({ + "FPS vs Num Assets": fig_fps, + "VRAM vs Num Assets": fig_mem, + "Num Assets Table": wandb.Table(dataframe=df), + }) + + +def compare_num_faces(): + csv_path = "outputs/benchmarks/ray_caster_benchmark_num_faces.csv" + df = pd.read_csv(csv_path) + + artifact = wandb.Artifact("ray_caster_benchmark_num_faces", type="dataset") + artifact.add_file(csv_path) + wandb.log_artifact(artifact) + + df["num_faces"] = df["num_faces"].astype(int) + df["avg_memory"] = df["avg_memory"].astype(float) + df["fps"] = 1.0 / (df["per_step_ms"] * 1e-3) + + fig_fps = px.line( + df, + x="num_faces", + y="fps", + markers=True, + title="FPS vs Num Faces", + ) + + fig_mem = px.line( + df, + x="num_faces", + y="avg_memory", + markers=True, + title="Average VRAM Usage (MB) vs Num Faces", + ) + + wandb.log({ + "FPS vs Num Faces": fig_fps, + "VRAM vs Num Faces": fig_mem, + "Num Faces Table": wandb.Table(dataframe=df), + }) + + +if __name__ == "__main__": + wandb.init(project="raycast-benchmarks", job_type="analysis", entity=os.getenv("WANDB_ENTITY", None)) + compare_single_vs_multi() + compare_num_assets_vram_cache() + compare_num_assets() + compare_num_faces() + wandb.finish() diff --git a/scripts/benchmarks/utils.py b/scripts/benchmarks/utils.py index ff2ca5c0114..026a3db08b4 100644 --- a/scripts/benchmarks/utils.py +++ b/scripts/benchmarks/utils.py @@ -102,3 +102,5 @@ def log_rl_policy_episode_lengths(benchmark: BaseIsaacBenchmark, value: list): # log max episode length measurement = SingleMeasurement(name="Max Episode Lengths", value=max(value), unit="float") benchmark.store_custom_measurement("train", measurement) + + diff --git a/scripts/demos/sensors/multi_mesh_raycaster.py b/scripts/demos/sensors/multi_mesh_raycaster.py new file mode 100644 index 00000000000..f31ece2c44d --- /dev/null +++ b/scripts/demos/sensors/multi_mesh_raycaster.py @@ -0,0 +1,290 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + + +"""Example on using the MultiMesh Raycaster sensor. + +Usage: + `python scripts/demos/sensors/multi_mesh_raycaster.py --num_envs 16 --asset_type ` +""" + +import argparse + +from isaaclab.app import AppLauncher + +# add argparse arguments +parser = argparse.ArgumentParser(description="Example on using the raycaster sensor.") +parser.add_argument("--num_envs", type=int, default=16, help="Number of environments to spawn.") +parser.add_argument( + "--asset_type", + type=str, + default="allegro_hand", + help="Asset type to use.", + choices=["allegro_hand", "anymal_d", "multi"], +) +# append AppLauncher cli args +AppLauncher.add_app_launcher_args(parser) +# parse the arguments +args_cli = parser.parse_args() + +# launch omniverse app +app_launcher = AppLauncher(args_cli) +simulation_app = app_launcher.app + +"""Rest everything follows.""" + +import random +import torch + +import omni.usd +from pxr import Gf, Sdf + +## +# Pre-defined configs +## +from isaaclab_assets.robots.allegro import ALLEGRO_HAND_CFG +from isaaclab_assets.robots.anymal import ANYMAL_D_CFG + +import isaaclab.sim as sim_utils +from isaaclab.assets import Articulation, AssetBaseCfg, RigidObjectCfg +from isaaclab.markers.config import VisualizationMarkersCfg +from isaaclab.scene import InteractiveScene, InteractiveSceneCfg +from isaaclab.sensors.ray_caster import MultiMeshRayCasterCfg, patterns +from isaaclab.utils import configclass +from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR + +RAY_CASTER_MARKER_CFG = VisualizationMarkersCfg( + markers={ + "hit": sim_utils.SphereCfg( + radius=0.01, + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(1.0, 0.0, 0.0)), + ), + }, +) + + +if args_cli.asset_type == "allegro_hand": + asset_cfg = ALLEGRO_HAND_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + ray_caster_cfg = MultiMeshRayCasterCfg( + prim_path="{ENV_REGEX_NS}/Robot", + update_period=1 / 60, + offset=MultiMeshRayCasterCfg.OffsetCfg(pos=(0, -0.1, 0.3)), + mesh_prim_paths=[ + "/World/Ground", + MultiMeshRayCasterCfg.RaycastTargetCfg(target_prim_expr="{ENV_REGEX_NS}/Robot/thumb_link_.*/visuals_xform"), + MultiMeshRayCasterCfg.RaycastTargetCfg(target_prim_expr="{ENV_REGEX_NS}/Robot/index_link.*/visuals_xform"), + MultiMeshRayCasterCfg.RaycastTargetCfg( + target_prim_expr="{ENV_REGEX_NS}/Robot/middle_link_.*/visuals_xform" + ), + MultiMeshRayCasterCfg.RaycastTargetCfg(target_prim_expr="{ENV_REGEX_NS}/Robot/ring_link_.*/visuals_xform"), + MultiMeshRayCasterCfg.RaycastTargetCfg(target_prim_expr="{ENV_REGEX_NS}/Robot/palm_link/visuals_xform"), + MultiMeshRayCasterCfg.RaycastTargetCfg(target_prim_expr="{ENV_REGEX_NS}/Robot/allegro_mount/visuals_xform"), + ], + ray_alignment="world", + pattern_cfg=patterns.GridPatternCfg(resolution=0.005, size=(0.4, 0.4), direction=(0, 0, -1)), + debug_vis=not args_cli.headless, + visualizer_cfg=RAY_CASTER_MARKER_CFG.replace(prim_path="/Visuals/RayCaster"), + ) + +elif args_cli.asset_type == "anymal_d": + asset_cfg = ANYMAL_D_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + ray_caster_cfg = MultiMeshRayCasterCfg( + prim_path="{ENV_REGEX_NS}/Robot", + update_period=1 / 60, + offset=MultiMeshRayCasterCfg.OffsetCfg(pos=(0, -0.1, 0.3)), + mesh_prim_paths=[ + "/World/Ground", + MultiMeshRayCasterCfg.RaycastTargetCfg(target_prim_expr="{ENV_REGEX_NS}/Robot/LF_.*/visuals"), + MultiMeshRayCasterCfg.RaycastTargetCfg(target_prim_expr="{ENV_REGEX_NS}/Robot/RF_.*/visuals"), + MultiMeshRayCasterCfg.RaycastTargetCfg(target_prim_expr="{ENV_REGEX_NS}/Robot/LH_.*/visuals"), + MultiMeshRayCasterCfg.RaycastTargetCfg(target_prim_expr="{ENV_REGEX_NS}/Robot/RH_.*/visuals"), + MultiMeshRayCasterCfg.RaycastTargetCfg(target_prim_expr="{ENV_REGEX_NS}/Robot/base/visuals"), + ], + ray_alignment="world", + pattern_cfg=patterns.GridPatternCfg(resolution=0.02, size=(2.5, 2.5), direction=(0, 0, -1)), + debug_vis=not args_cli.headless, + visualizer_cfg=RAY_CASTER_MARKER_CFG.replace(prim_path="/Visuals/RayCaster"), + ) + +elif args_cli.asset_type == "multi": + asset_cfg = RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/Object", + spawn=sim_utils.MultiAssetSpawnerCfg( + assets_cfg=[ + sim_utils.CuboidCfg( + size=(0.3, 0.3, 0.3), + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(1.0, 0.0, 0.0), metallic=0.2), + ), + sim_utils.SphereCfg( + radius=0.3, + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.0, 0.0, 1.0), metallic=0.2), + ), + ], + random_choice=True, + rigid_props=sim_utils.RigidBodyPropertiesCfg( + solver_position_iteration_count=4, solver_velocity_iteration_count=0 + ), + mass_props=sim_utils.MassPropertiesCfg(mass=1.0), + collision_props=sim_utils.CollisionPropertiesCfg(), + ), + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, 2.0)), + ) + ray_caster_cfg = MultiMeshRayCasterCfg( + prim_path="{ENV_REGEX_NS}/Object", + update_period=1 / 60, + offset=MultiMeshRayCasterCfg.OffsetCfg(pos=(0, 0.0, 0.6)), + mesh_prim_paths=[ + "/World/Ground", + MultiMeshRayCasterCfg.RaycastTargetCfg(target_prim_expr="{ENV_REGEX_NS}/Object"), + ], + ray_alignment="world", + pattern_cfg=patterns.GridPatternCfg(resolution=0.01, size=(0.6, 0.6), direction=(0, 0, -1)), + debug_vis=not args_cli.headless, + visualizer_cfg=RAY_CASTER_MARKER_CFG.replace(prim_path="/Visuals/RayCaster"), + ) +else: + raise ValueError(f"Unknown asset type: {args_cli.asset_type}") + + +@configclass +class RaycasterSensorSceneCfg(InteractiveSceneCfg): + """Design the scene with sensors on the asset.""" + + # ground plane + ground = AssetBaseCfg( + prim_path="/World/Ground", + spawn=sim_utils.UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Environments/Terrains/rough_plane.usd", + scale=(1, 1, 1), + ), + ) + + # lights + dome_light = AssetBaseCfg( + prim_path="/World/Light", spawn=sim_utils.DomeLightCfg(intensity=3000.0, color=(0.75, 0.75, 0.75)) + ) + + # asset + asset = asset_cfg + # ray caster + ray_caster = ray_caster_cfg + + +def randomize_shape_color(prim_path_expr: str): + """Randomize the color of the geometry.""" + + # acquire stage + stage = omni.usd.get_context().get_stage() + # resolve prim paths for spawning and cloning + prim_paths = sim_utils.find_matching_prim_paths(prim_path_expr) + # manually clone prims if the source prim path is a regex expression + + with Sdf.ChangeBlock(): + for prim_path in prim_paths: + print("Applying prim scale to:", prim_path) + # spawn single instance + prim_spec = Sdf.CreatePrimInLayer(stage.GetRootLayer(), prim_path) + + # DO YOUR OWN OTHER KIND OF RANDOMIZATION HERE! + # Note: Just need to acquire the right attribute about the property you want to set + # Here is an example on setting color randomly + color_spec = prim_spec.GetAttributeAtPath(prim_path + "/geometry/material/Shader.inputs:diffuseColor") + color_spec.default = Gf.Vec3f(random.random(), random.random(), random.random()) + + # randomize scale + scale_spec = prim_spec.GetAttributeAtPath(prim_path + ".xformOp:scale") + scale_spec.default = Gf.Vec3f(random.uniform(0.5, 1.5), random.uniform(0.5, 1.5), random.uniform(0.5, 1.5)) + + +def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): + """Run the simulator.""" + # Define simulation stepping + sim_dt = sim.get_physics_dt() + sim_time = 0.0 + count = 0 + + triggered = True + countdown = 42 + + # Simulate physics + while simulation_app.is_running(): + + if count % 500 == 0: + # reset counter + count = 0 + # reset the scene entities + # root state + root_state = scene["asset"].data.default_root_state.clone() + root_state[:, :3] += scene.env_origins + scene["asset"].write_root_pose_to_sim(root_state[:, :7]) + scene["asset"].write_root_velocity_to_sim(root_state[:, 7:]) + + if isinstance(scene["asset"], Articulation): + # set joint positions with some noise + joint_pos, joint_vel = ( + scene["asset"].data.default_joint_pos.clone(), + scene["asset"].data.default_joint_vel.clone(), + ) + joint_pos += torch.rand_like(joint_pos) * 0.1 + scene["asset"].write_joint_state_to_sim(joint_pos, joint_vel) + # clear internal buffers + scene.reset() + print("[INFO]: Resetting Asset state...") + + if isinstance(scene["asset"], Articulation): + # -- generate actions/commands + targets = scene["asset"].data.default_joint_pos + 5 * ( + torch.rand_like(scene["asset"].data.default_joint_pos) - 0.5 + ) + # -- apply action to the asset + scene["asset"].set_joint_position_target(targets) + # -- write data to sim + scene.write_data_to_sim() + # perform step + sim.step() + # update sim-time + sim_time += sim_dt + count += 1 + # update buffers + scene.update(sim_dt) + + if not triggered: + if countdown > 0: + countdown -= 1 + continue + data = scene["ray_caster"].data.ray_hits_w.cpu().numpy() # noqa: F841 + triggered = True + else: + continue + + +def main(): + """Main function.""" + + # Initialize the simulation context + sim_cfg = sim_utils.SimulationCfg(dt=0.005, device=args_cli.device) + sim = sim_utils.SimulationContext(sim_cfg) + # Set main camera + sim.set_camera_view(eye=[3.5, 3.5, 3.5], target=[0.0, 0.0, 0.0]) + # design scene + scene_cfg = RaycasterSensorSceneCfg(num_envs=args_cli.num_envs, env_spacing=2.0, replicate_physics=False) + scene = InteractiveScene(scene_cfg) + + if args_cli.asset_type == "multi": + randomize_shape_color(scene_cfg.asset.prim_path.format(ENV_REGEX_NS="/World/envs/env_.*")) + + # Play the simulator + sim.reset() + # Now we are ready! + print("[INFO]: Setup complete...") + # Run the simulator + run_simulator(sim, scene) + + +if __name__ == "__main__": + # run the main function + main() + # close sim app + simulation_app.close() diff --git a/source/isaaclab/config/extension.toml b/source/isaaclab/config/extension.toml index fde6cb9bf91..2c62e077cd1 100644 --- a/source/isaaclab/config/extension.toml +++ b/source/isaaclab/config/extension.toml @@ -1,7 +1,7 @@ [package] # Note: Semantic Versioning is used: https://semver.org/ -version = "0.46.2" +version = "0.46.3" # Description title = "Isaac Lab framework for Robot Learning" @@ -17,7 +17,7 @@ requirements = [ "prettytable==3.3.0", "toml", "hidapi", - "gymnasium==0.29.0", + "gymnasium==1.2.0", "trimesh" ] diff --git a/source/isaaclab/docs/CHANGELOG.rst b/source/isaaclab/docs/CHANGELOG.rst index c33e307eebe..692b8776d89 100644 --- a/source/isaaclab/docs/CHANGELOG.rst +++ b/source/isaaclab/docs/CHANGELOG.rst @@ -1,6 +1,21 @@ Changelog --------- +0.46.3 (2025-09-06) +~~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Added :meth:`~isaaclab.sim.utils.resolve_prim_pose` to resolve the pose of a prim with respect to another prim. +* Added :meth:`~isaaclab.sim.utils.resolve_prim_scale` to resolve the scale of a prim in the world frame. + +Changed +^^^^^^^ + +* Added parsing of instanced prims in :meth:`~isaaclab.sim.utils.get_all_matching_child_prims` and :meth:`~isaaclab.sim.utils.get_first_matching_child_prim`. + Earlier, instanced prims were skipped since :meth:`Usd.Prim.GetChildren` does not return instanced prims. + 0.46.2 (2025-09-13) ~~~~~~~~~~~~~~~~~~~ diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/__init__.py b/source/isaaclab/isaaclab/sensors/ray_caster/__init__.py index a4fe1bce519..bd87b8d59b3 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/__init__.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/__init__.py @@ -6,6 +6,10 @@ """Sub-module for Warp-based ray-cast sensor.""" from . import patterns +from .multi_mesh_ray_caster import MultiMeshRayCaster +from .multi_mesh_ray_caster_camera import MultiMeshRayCasterCamera +from .multi_mesh_ray_caster_camera_cfg import MultiMeshRayCasterCameraCfg +from .multi_mesh_ray_caster_cfg import MultiMeshRayCasterCfg from .ray_caster import RayCaster from .ray_caster_camera import RayCasterCamera from .ray_caster_camera_cfg import RayCasterCameraCfg diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster.py b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster.py new file mode 100644 index 00000000000..6536f0fcf89 --- /dev/null +++ b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster.py @@ -0,0 +1,415 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +"""Multi-mesh ray casting sensor implementation. + +This file adds support for ray casting against multiple (possibly regex-selected) mesh targets. +""" + +import numpy as np +import re +import torch +import trimesh +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import carb +import omni.log +import warp as wp +from isaacsim.core.prims import XFormPrim +from pxr import UsdPhysics + +import isaaclab.sim as sim_utils +from isaaclab.utils.math import matrix_from_quat, quat_mul +from isaaclab.utils.mesh import PRIMITIVE_MESH_TYPES, create_trimesh_from_geom_mesh, create_trimesh_from_geom_shape +from isaaclab.utils.warp import convert_to_warp_mesh, raycast_dynamic_meshes + +from .multi_mesh_ray_caster_data import MultiMeshRayCasterData +from .ray_caster import RayCaster + +if TYPE_CHECKING: + from .multi_mesh_ray_caster_cfg import MultiMeshRayCasterCfg + + +class MultiMeshRayCaster(RayCaster): + """A multi-mesh ray-casting sensor. + + The ray-caster uses a set of rays to detect collisions with meshes in the scene. The rays are + defined in the sensor's local coordinate frame. The sensor can be configured to ray-cast against + a set of meshes with a given ray pattern. + + The meshes are parsed from the list of primitive paths provided in the configuration. These are then + converted to warp meshes and stored in the :attr:`warp_meshes` list. The ray-caster then ray-casts against + these warp meshes using the ray pattern provided in the configuration. + + Compared to the default RayCaster, the MultiMeshRayCaster provides additional functionality and flexibility as + an extension of the default RayCaster with the following enhancements: + + - Raycasting against multiple target types : Supports primitive shapes (spheres, cubes, …) as well as arbitrary + meshes. + - Dynamic mesh tracking : Keeps track of specified meshes, enabling raycasting against moving parts + (e.g., robot links, articulated bodies, or dynamic obstacles). + - Memory-efficient caching : Avoids redundant memory usage by reusing mesh data across environments. + + Example usage to raycast against the visual meshes of a robot (e.g. anymal): + .. code-block:: python + ray_caster_cfg = MultiMeshRayCasterCfg( + prim_path="{ENV_REGEX_NS}/Robot", + mesh_prim_paths=[ + "/World/Ground", + MultiMeshRayCasterCfg.RaycastTargetCfg(target_prim_expr="{ENV_REGEX_NS}/Robot/LF_.*/visuals"), + MultiMeshRayCasterCfg.RaycastTargetCfg(target_prim_expr="{ENV_REGEX_NS}/Robot/RF_.*/visuals"), + MultiMeshRayCasterCfg.RaycastTargetCfg(target_prim_expr="{ENV_REGEX_NS}/Robot/LH_.*/visuals"), + MultiMeshRayCasterCfg.RaycastTargetCfg(target_prim_expr="{ENV_REGEX_NS}/Robot/RH_.*/visuals"), + MultiMeshRayCasterCfg.RaycastTargetCfg(target_prim_expr="{ENV_REGEX_NS}/Robot/base/visuals"), + ], + ray_alignment="world", + pattern_cfg=patterns.GridPatternCfg(resolution=0.02, size=(2.5, 2.5), direction=(0, 0, -1)), + ) + """ + + cfg: MultiMeshRayCasterCfg + """The configuration parameters.""" + + mesh_offsets: dict[str, tuple[torch.Tensor, torch.Tensor]] = {} + + def __init__(self, cfg: MultiMeshRayCasterCfg): + """Initializes the ray-caster object. + + Args: + cfg: The configuration parameters. + """ + # Initialize base class + super().__init__(cfg) + + # Create empty variables for storing output data + self._num_meshes_per_env: dict[str, int] = {} + """Keeps track of the number of meshes per env for each ray_cast target. + Since we allow regex indexing (e.g. env_*/object_*) they can differ + """ + + self._raycast_targets_cfg: list[MultiMeshRayCasterCfg.RaycastTargetCfg] = [] + for target in self.cfg.mesh_prim_paths: + # Legacy support for string targets. Treat them as global targets. + if isinstance(target, str): + self._raycast_targets_cfg.append( + cfg.RaycastTargetCfg(target_prim_expr=target, track_mesh_transforms=False) + ) + else: + self._raycast_targets_cfg.append(target) + + # Resolve regex namespace if set + for cfg in self._raycast_targets_cfg: + cfg.target_prim_expr = cfg.target_prim_expr.format(ENV_REGEX_NS="/World/envs/env_.*") + + # overwrite the data class + self._data = MultiMeshRayCasterData() + + def __str__(self) -> str: + """Returns: A string containing information about the instance.""" + + return ( + f"Ray-caster @ '{self.cfg.prim_path}': \n" + f"\tview type : {self._view.__class__}\n" + f"\tupdate period (s) : {self.cfg.update_period}\n" + f"\tnumber of meshes : {self._num_envs} x {sum(self._num_meshes_per_env.values())} \n" + f"\tnumber of sensors : {self._view.count}\n" + f"\tnumber of rays/sensor: {self.num_rays}\n" + f"\ttotal number of rays : {self.num_rays * self._view.count}" + ) + + """ + Properties + """ + + @property + def data(self) -> MultiMeshRayCasterData: + # update sensors if needed + self._update_outdated_buffers() + # return the data + return self._data + + """ + Implementation. + """ + + def _initialize_warp_meshes(self): + """Parse mesh prim expressions, build (or reuse) Warp meshes, and cache per-env mesh IDs. + + High-level steps (per target expression): + 1. Resolve matching prims by regex/path expression. + 2. Collect supported mesh child prims; merge into a single mesh if configured. + 3. Deduplicate identical vertex buffers (exact match) to avoid uploading duplicates to Warp. + 4. Partition mesh IDs per environment or mark as globally shared. + 5. Optionally create physics views (articulation / rigid body / fallback XForm) and cache local offsets. + + Exceptions: + Raises a RuntimeError if: + - No prims match the provided expression. + - No supported mesh prims are found under a matched prim. + - Multiple mesh prims are found but merging is disabled. + """ + multi_mesh_ids: dict[str, list[list[int]]] = {} + for target_cfg in self._raycast_targets_cfg: + # target prim path to ray cast against + target_prim_path = target_cfg.target_prim_expr + # check if mesh already casted into warp mesh and get the number of meshes per env + if target_prim_path in multi_mesh_ids: + self._num_meshes_per_env[target_prim_path] = len(multi_mesh_ids[target_prim_path]) // self._num_envs + continue + + # find all matching prim paths to provided expression of the target + target_prims = sim_utils.find_matching_prims(target_prim_path) + if len(target_prims) == 0: + raise RuntimeError(f"Failed to find a prim at path expression: {target_prim_path}") + + is_global_prim = ( + len(target_prims) == 1 + ) # If only one prim is found, treat it as a global prim. Either it's a single global object (e.g. ground) or we are only using one env. + + loaded_vertices: list[np.ndarray | None] = [] + wp_mesh_ids = [] + + for target_prim in target_prims: + # Reuse previously parsed shared mesh instance if possible. + if target_cfg.is_shared and len(wp_mesh_ids) > 0: + # Verify if this mesh has already been registered in an earlier environment. + # Note, this check may fail, if the prim path is not following the env_.* pattern + # Which (worst case) leads to parsing the mesh and skipping registering it at a later stage + curr_prim_base_path = re.sub(r"env_\d+", "env_0", str(target_prim.GetPath())) # + if curr_prim_base_path in MultiMeshRayCaster.meshes: + MultiMeshRayCaster.meshes[str(target_prim.GetPath())] = MultiMeshRayCaster.meshes[ + curr_prim_base_path + ] + # Reuse mesh imported by another ray-cast sensor (global cache). + if str(target_prim.GetPath()) in MultiMeshRayCaster.meshes: + wp_mesh_ids.append(MultiMeshRayCaster.meshes[str(target_prim.GetPath())].id) + loaded_vertices.append(None) + continue + + mesh_prims = sim_utils.get_all_matching_child_prims( + target_prim.GetPath(), lambda prim: prim.GetTypeName() in PRIMITIVE_MESH_TYPES + ["Mesh"] + ) + if len(mesh_prims) == 0: + warn_msg = ( + f"No mesh prims found at path: {target_prim.GetPath()} with supported types:" + f" {PRIMITIVE_MESH_TYPES + ['Mesh']}" + " Skipping this target." + ) + for prim in sim_utils.get_all_matching_child_prims(target_prim.GetPath(), lambda prim: True): + warn_msg += f"\n - Available prim '{prim.GetPath()}' of type '{prim.GetTypeName()}'" + carb.log_warn(warn_msg) + continue + + trimesh_meshes = [] + + for mesh_prim in mesh_prims: + # check if valid + if mesh_prim is None or not mesh_prim.IsValid(): + raise RuntimeError(f"Invalid mesh prim path: {target_prim}") + + if mesh_prim.GetTypeName() == "Mesh": + mesh = create_trimesh_from_geom_mesh(mesh_prim) + else: + mesh = create_trimesh_from_geom_shape(mesh_prim) + scale = sim_utils.resolve_prim_scale(mesh_prim) + mesh.apply_scale(scale) + + relative_pos, relative_quat = sim_utils.resolve_prim_pose(mesh_prim, target_prim) + relative_pos = torch.tensor(relative_pos, dtype=torch.float32) + relative_quat = torch.tensor(relative_quat, dtype=torch.float32) + + rotation = matrix_from_quat(relative_quat) + transform = np.eye(4) + transform[:3, :3] = rotation.numpy() + transform[:3, 3] = relative_pos.numpy() + mesh.apply_transform(transform) + + # add to list of parsed meshes + trimesh_meshes.append(mesh) + + if len(trimesh_meshes) == 1: + trimesh_mesh = trimesh_meshes[0] + elif target_cfg.merge_prim_meshes: + # combine all trimesh meshes into a single mesh + trimesh_mesh = trimesh.util.concatenate(trimesh_meshes) + else: + raise RuntimeError( + f"Multiple mesh prims found at path: {target_prim.GetPath()} but merging is disabled. Please" + " enable `merge_prim_meshes` in the configuration or specify each mesh separately." + ) + + # check if the mesh is already registered, if so only reference the mesh + registered_idx = _registered_points_idx(trimesh_mesh.vertices, loaded_vertices) + if registered_idx != -1 and self.cfg.reference_meshes: + omni.log.info("Found a duplicate mesh, only reference the mesh.") + # Found a duplicate mesh, only reference the mesh. + loaded_vertices.append(None) + wp_mesh_ids.append(wp_mesh_ids[registered_idx]) + else: + loaded_vertices.append(trimesh_mesh.vertices) + wp_mesh = convert_to_warp_mesh(trimesh_mesh.vertices, trimesh_mesh.faces, device=self.device) + MultiMeshRayCaster.meshes[str(target_prim.GetPath())] = wp_mesh + wp_mesh_ids.append(wp_mesh.id) + + # print info + if registered_idx != -1: + omni.log.info(f"Found duplicate mesh for mesh prims under path '{target_prim.GetPath()}'.") + else: + omni.log.info( + f"Read '{len(mesh_prims)}' mesh prims under path '{target_prim.GetPath()}' with" + f" {len(trimesh_mesh.vertices)} vertices and {len(trimesh_mesh.faces)} faces." + ) + + if is_global_prim: + # reference the mesh for each environment to ray cast against + multi_mesh_ids[target_prim_path] = [wp_mesh_ids] * self._num_envs + self._num_meshes_per_env[target_prim_path] = len(wp_mesh_ids) + else: + # split up the meshes for each environment. Little bit ugly, since + # the current order is interleaved (env1_obj1, env1_obj2, env2_obj1, env2_obj2, ...) + multi_mesh_ids[target_prim_path] = [] + mesh_idx = 0 + n_meshes_per_env = len(wp_mesh_ids) // self._num_envs + self._num_meshes_per_env[target_prim_path] = n_meshes_per_env + for _ in range(self._num_envs): + multi_mesh_ids[target_prim_path].append(wp_mesh_ids[mesh_idx : mesh_idx + n_meshes_per_env]) + mesh_idx += n_meshes_per_env + + if target_cfg.track_mesh_transforms: + mesh_prim = sim_utils.find_first_matching_prim(target_prim_path) + self.mesh_views[target_prim_path], MultiMeshRayCaster.mesh_offsets[target_prim_path] = ( + self._get_trackable_prim_view(target_prim_path) + ) + + # throw an error if no meshes are found + if all([target_cfg.target_prim_expr not in multi_mesh_ids for target_cfg in self._raycast_targets_cfg]): + raise RuntimeError( + f"No meshes found for ray-casting! Please check the mesh prim paths: {self.cfg.mesh_prim_paths}" + ) + + total_n_meshes_per_env = sum(self._num_meshes_per_env.values()) + self._mesh_positions_w = torch.zeros(self._num_envs, total_n_meshes_per_env, 3, device=self.device) + self._mesh_orientations_w = torch.zeros(self._num_envs, total_n_meshes_per_env, 4, device=self.device) + + # Update the mesh positions and rotations + mesh_idx = 0 + for target_cfg in self._raycast_targets_cfg: + n_meshes = self._num_meshes_per_env[target_cfg.target_prim_expr] + + # update position of the target meshes + pos_w, ori_w = [], [] + for prim in sim_utils.find_matching_prims(target_cfg.target_prim_expr): + translation, quat = sim_utils.resolve_prim_pose(prim) + pos_w.append(translation) + ori_w.append(quat) + pos_w = torch.tensor(pos_w, device=self.device, dtype=torch.float32).view(-1, n_meshes, 3) + ori_w = torch.tensor(ori_w, device=self.device, dtype=torch.float32).view(-1, n_meshes, 4) + + self._mesh_positions_w[:, mesh_idx : mesh_idx + n_meshes] = pos_w + self._mesh_orientations_w[:, mesh_idx : mesh_idx + n_meshes] = ori_w + mesh_idx += n_meshes + + # flatten the list of meshes that are included in mesh_prim_paths of the specific ray caster + multi_mesh_ids_flattened = [] + for env_idx in range(self._num_envs): + meshes_in_env = [] + for target_cfg in self._raycast_targets_cfg: + meshes_in_env.extend(multi_mesh_ids[target_cfg.target_prim_expr][env_idx]) + multi_mesh_ids_flattened.append(meshes_in_env) + + self._mesh_views = [ + self.mesh_views[target_cfg.target_prim_expr] if target_cfg.track_mesh_transforms else None + for target_cfg in self._raycast_targets_cfg + ] + + # save a warp array with mesh ids that is passed to the raycast function + self._mesh_ids_wp = wp.array2d(multi_mesh_ids_flattened, dtype=wp.uint64, device=self.device) + + def _initialize_rays_impl(self): + super()._initialize_rays_impl() + if self.cfg.update_mesh_ids: + self._data.ray_mesh_ids = torch.zeros( + self._num_envs, self.num_rays, 1, device=self.device, dtype=torch.int16 + ) + + def _update_buffers_impl(self, env_ids: Sequence[int]): + """Fills the buffers of the sensor data. + + Args: + env_ids: The environment ids to update. + """ + + self._update_ray_infos(env_ids) + + # Update the mesh positions and rotations + mesh_idx = 0 + for view, target_cfg in zip(self._mesh_views, self._raycast_targets_cfg): + if not target_cfg.track_mesh_transforms: + mesh_idx += self._num_meshes_per_env[target_cfg.target_prim_expr] + continue + + # update position of the target meshes + pos_w, ori_w = sim_utils.obtain_world_pose_from_view(view, None) + pos_w = pos_w.squeeze(0) if len(pos_w.shape) == 3 else pos_w + ori_w = ori_w.squeeze(0) if len(ori_w.shape) == 3 else ori_w + + if target_cfg.target_prim_expr in MultiMeshRayCaster.mesh_offsets: + pos_offset, ori_offset = MultiMeshRayCaster.mesh_offsets[target_cfg.target_prim_expr] + pos_w -= pos_offset + ori_w = quat_mul(ori_offset.expand(ori_w.shape[0], -1), ori_w) + + count = view.count + if count != 1: # Mesh is not global, i.e. we have different meshes for each env + count = count // self._num_envs + pos_w = pos_w.view(self._num_envs, count, 3) + ori_w = ori_w.view(self._num_envs, count, 4) + + self._mesh_positions_w[:, mesh_idx : mesh_idx + count] = pos_w + self._mesh_orientations_w[:, mesh_idx : mesh_idx + count] = ori_w + mesh_idx += count + + self._data.ray_hits_w[env_ids], _, _, _, mesh_ids = raycast_dynamic_meshes( + self._ray_starts_w[env_ids], + self._ray_directions_w[env_ids], + mesh_ids_wp=self._mesh_ids_wp, # list with shape num_envs x num_meshes_per_env + max_dist=self.cfg.max_distance, + mesh_positions_w=self._mesh_positions_w[env_ids], + mesh_orientations_w=self._mesh_orientations_w[env_ids], + return_mesh_id=self.cfg.update_mesh_ids, + ) + + if self.cfg.update_mesh_ids: + self._data.ray_mesh_ids[env_ids] = mesh_ids + + def __del__(self): + super().__del__() + if RayCaster._instance_count == 0: + MultiMeshRayCaster.mesh_offsets.clear() + + +""" +Helper functions +""" + + +def _registered_points_idx(points: np.ndarray, registered_points: list[np.ndarray | None]) -> int: + """Check if the points are already registered in the list of registered points. + + Args: + points: The points to check. + registered_points: The list of registered points. + + Returns: + The index of the registered points if found, otherwise -1. + """ + for idx, reg_points in enumerate(registered_points): + if reg_points is None: + continue + if reg_points.shape == points.shape and (reg_points == points).all(): + return idx + return -1 diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_camera.py b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_camera.py new file mode 100644 index 00000000000..2f4c396e197 --- /dev/null +++ b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_camera.py @@ -0,0 +1,212 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import isaaclab.sim as sim_utils +import isaaclab.utils.math as math_utils +from isaaclab.utils.warp import raycast_dynamic_meshes + +from .multi_mesh_ray_caster import MultiMeshRayCaster +from .multi_mesh_ray_caster_camera_data import MultiMeshRayCasterCameraData +from .ray_caster_camera import RayCasterCamera + +if TYPE_CHECKING: + from .multi_mesh_ray_caster_camera_cfg import MultiMeshRayCasterCameraCfg + + +class MultiMeshRayCasterCamera(RayCasterCamera, MultiMeshRayCaster): + """A multi-mesh ray-casting camera sensor. + + The ray-caster camera uses a set of rays to get the distances to meshes in the scene. The rays are + defined in the sensor's local coordinate frame. The sensor has the same interface as the + :class:`isaaclab.sensors.Camera` that implements the camera class through USD camera prims. + However, this class provides a faster image generation. The sensor converts meshes from the list of + primitive paths provided in the configuration to Warp meshes. The camera then ray-casts against these + Warp meshes only. + + Currently, only the following annotators are supported: + + - ``"distance_to_camera"``: An image containing the distance to camera optical center. + - ``"distance_to_image_plane"``: An image containing distances of 3D points from camera plane along camera's z-axis. + - ``"normals"``: An image containing the local surface normal vectors at each pixel. + """ + + cfg: MultiMeshRayCasterCameraCfg + """The configuration parameters.""" + + def __init__(self, cfg: MultiMeshRayCasterCameraCfg): + """Initializes the camera object. + + Args: + cfg: The configuration parameters. + + Raises: + ValueError: If the provided data types are not supported by the ray-caster camera. + """ + self._check_supported_data_types(cfg) + # initialize base class + MultiMeshRayCaster.__init__(self, cfg) + # create empty variables for storing output data + self._data = MultiMeshRayCasterCameraData() + + def __str__(self) -> str: + """Returns: A string containing information about the instance.""" + return ( + f"Multi-Mesh Ray-Caster-Camera @ '{self.cfg.prim_path}': \n" + f"\tview type : {self._view.__class__}\n" + f"\tupdate period (s) : {self.cfg.update_period}\n" + f"\tnumber of meshes : {len(MultiMeshRayCaster.meshes)}\n" + f"\tnumber of sensors : {self._view.count}\n" + f"\tnumber of rays/sensor: {self.num_rays}\n" + f"\ttotal number of rays : {self.num_rays * self._view.count}\n" + f"\timage shape : {self.image_shape}" + ) + + """ + Implementation. + """ + + def _initialize_warp_meshes(self): + MultiMeshRayCaster._initialize_warp_meshes(self) + + def _create_buffers(self): + super()._create_buffers() + self._data.image_mesh_ids = torch.zeros( + self._num_envs, *self.image_shape, 1, device=self.device, dtype=torch.int16 + ) + + def _initialize_rays_impl(self): + # Create all indices buffer + self._ALL_INDICES = torch.arange(self._view.count, device=self._device, dtype=torch.long) + # Create frame count buffer + self._frame = torch.zeros(self._view.count, device=self._device, dtype=torch.long) + # create buffers + self._create_buffers() + # compute intrinsic matrices + self._compute_intrinsic_matrices() + # compute ray stars and directions + self.ray_starts, self.ray_directions = self.cfg.pattern_cfg.func( + self.cfg.pattern_cfg, self._data.intrinsic_matrices, self._device + ) + self.num_rays = self.ray_directions.shape[1] + # create buffer to store ray hits + self.ray_hits_w = torch.zeros(self._view.count, self.num_rays, 3, device=self._device) + # set offsets + quat_w = math_utils.convert_camera_frame_orientation_convention( + torch.tensor([self.cfg.offset.rot], device=self._device), origin=self.cfg.offset.convention, target="world" + ) + self._offset_quat = quat_w.repeat(self._view.count, 1) + self._offset_pos = torch.tensor(list(self.cfg.offset.pos), device=self._device).repeat(self._view.count, 1) + + self._data.quat_w = torch.zeros(self._view.count, 4, device=self.device) + self._data.pos_w = torch.zeros(self._view.count, 3, device=self.device) + + self._ray_starts_w = torch.zeros(self._view.count, self.num_rays, 3, device=self.device) + self._ray_directions_w = torch.zeros(self._view.count, self.num_rays, 3, device=self.device) + + def _update_ray_infos(self, env_ids: Sequence[int]): + """Updates the ray information buffers.""" + + # compute poses from current view + pos_w, quat_w = sim_utils.obtain_world_pose_from_view(self._view, env_ids) + pos_w, quat_w = math_utils.combine_frame_transforms( + pos_w, quat_w, self._offset_pos[env_ids], self._offset_quat[env_ids] + ) + # update the data + self._data.pos_w[env_ids] = pos_w + self._data.quat_w_world[env_ids] = quat_w + + # note: full orientation is considered + ray_starts_w = math_utils.quat_apply(quat_w.repeat(1, self.num_rays), self.ray_starts[env_ids]) + ray_starts_w += pos_w.unsqueeze(1) + ray_directions_w = math_utils.quat_apply(quat_w.repeat(1, self.num_rays), self.ray_directions[env_ids]) + + self._ray_starts_w[env_ids] = ray_starts_w + self._ray_directions_w[env_ids] = ray_directions_w + + def _update_buffers_impl(self, env_ids: Sequence[int] | torch.Tensor | None): + """Fills the buffers of the sensor data.""" + self._update_ray_infos(env_ids) + + # increment frame count + if env_ids is None: + env_ids_tensor = torch.arange(self._num_envs, device=self.device) + elif not isinstance(env_ids, torch.Tensor): + env_ids_tensor = torch.tensor(env_ids, device=self.device) + else: + env_ids_tensor = env_ids + + self._frame[env_ids_tensor] += 1 + + # Update the mesh positions and rotations + mesh_idx = 0 + for view, target_cfg in zip(self._mesh_views, self._raycast_targets_cfg): + if not target_cfg.track_mesh_transforms: + mesh_idx += self._num_meshes_per_env[target_cfg.target_prim_expr] + continue + + # update position of the target meshes + pos_w, ori_w = sim_utils.obtain_world_pose_from_view(view, None) + pos_w = pos_w.squeeze(0) if len(pos_w.shape) == 3 else pos_w + ori_w = ori_w.squeeze(0) if len(ori_w.shape) == 3 else ori_w + + if target_cfg.target_prim_expr in MultiMeshRayCaster.mesh_offsets: + pos_offset, ori_offset = MultiMeshRayCaster.mesh_offsets[target_cfg.target_prim_expr] + pos_w -= pos_offset + ori_w = math_utils.quat_mul(ori_offset.expand(ori_w.shape[0], -1), ori_w) + + count = view.count + if not target_cfg.is_global: + count = count // self._num_envs + pos_w = pos_w.view(self._num_envs, count, 3) + ori_w = ori_w.view(self._num_envs, count, 4) + + self._mesh_positions_w[:, mesh_idx : mesh_idx + count] = pos_w + self._mesh_orientations_w[:, mesh_idx : mesh_idx + count] = ori_w + mesh_idx += count + + # ray cast and store the hits + self.ray_hits_w, ray_depth, ray_normal, _, ray_mesh_ids = raycast_dynamic_meshes( + self._ray_starts_w[env_ids], + self._ray_directions_w[env_ids], + mesh_ids_wp=self._mesh_ids_wp, # list with shape num_envs x num_meshes_per_env + max_dist=self.cfg.max_distance, + mesh_positions_w=self._mesh_positions_w[env_ids], + mesh_orientations_w=self._mesh_orientations_w[env_ids], + return_distance=any( + [name in self.cfg.data_types for name in ["distance_to_image_plane", "distance_to_camera"]] + ), + return_normal="normals" in self.cfg.data_types, + return_mesh_id=self.cfg.update_mesh_ids, + ) + + # update output buffers + if "distance_to_image_plane" in self.cfg.data_types: + # note: data is in camera frame so we only take the first component (z-axis of camera frame) + distance_to_image_plane = ( + math_utils.quat_apply( + math_utils.quat_inv(self._data.quat_w_world[env_ids]).repeat(1, self.num_rays), + (ray_depth[:, :, None] * self._ray_directions_w[env_ids]), + ) + )[:, :, 0] + # apply the maximum distance after the transformation + distance_to_image_plane = torch.clip(distance_to_image_plane, max=self.cfg.max_distance) + self._data.output["distance_to_image_plane"][env_ids_tensor] = distance_to_image_plane.view( + -1, *self.image_shape, 1 + ) + if "distance_to_camera" in self.cfg.data_types: + self._data.output["distance_to_camera"][env_ids_tensor] = torch.clip( + ray_depth.view(-1, *self.image_shape, 1), max=self.cfg.max_distance + ) + if "normals" in self.cfg.data_types: + self._data.output["normals"][env_ids_tensor] = ray_normal.view(-1, *self.image_shape, 3) + + if self.cfg.update_mesh_ids: + self._data.image_mesh_ids[env_ids_tensor] = ray_mesh_ids.view(-1, *self.image_shape, 1) diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_camera_cfg.py b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_camera_cfg.py new file mode 100644 index 00000000000..82d62482032 --- /dev/null +++ b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_camera_cfg.py @@ -0,0 +1,19 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Configuration for the ray-cast camera sensor.""" + +from isaaclab.utils import configclass + +from .multi_mesh_ray_caster_camera import MultiMeshRayCasterCamera +from .multi_mesh_ray_caster_cfg import MultiMeshRayCasterCfg +from .ray_caster_camera_cfg import RayCasterCameraCfg + + +@configclass +class MultiMeshRayCasterCameraCfg(RayCasterCameraCfg, MultiMeshRayCasterCfg): + """Configuration for the multi-mesh ray-cast camera sensor.""" + + class_type: type = MultiMeshRayCasterCamera diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_camera_data.py b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_camera_data.py new file mode 100644 index 00000000000..f8c7aea8863 --- /dev/null +++ b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_camera_data.py @@ -0,0 +1,21 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Configuration for the ray-cast sensor.""" +import torch + +from isaaclab.sensors.camera import CameraData + +from .ray_caster_data import RayCasterData + + +class MultiMeshRayCasterCameraData(CameraData, RayCasterData): + """Data container for the multi-mesh ray-cast sensor.""" + + image_mesh_ids: torch.Tensor = None + """The mesh ids of the image pixels. + + Shape is (N, H, W, 1), where N is the number of sensors, H and W are the height and width of the image, + and 1 is the number of mesh ids per pixel.""" diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_cfg.py b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_cfg.py new file mode 100644 index 00000000000..238eba97eb3 --- /dev/null +++ b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_cfg.py @@ -0,0 +1,70 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + + +"""Configuration for the ray-cast sensor.""" + + +from dataclasses import MISSING + +from isaaclab.utils import configclass + +from .multi_mesh_ray_caster import MultiMeshRayCaster +from .ray_caster_cfg import RayCasterCfg + + +@configclass +class MultiMeshRayCasterCfg(RayCasterCfg): + """Configuration for the multi-mesh ray-cast sensor.""" + + @configclass + class RaycastTargetCfg: + """Configuration for different ray-cast targets.""" + + target_prim_expr: str = MISSING + """The regex to specify the target prim to ray cast against.""" + + is_shared: bool = False + """Whether the target prim is assumed to be the same mesh across all environments. Defaults to False. + + If True, only the first mesh is read and then reused for all environments, rather than re-parsed. + This provides a startup performance boost when there are many environments that all use the same asset. + + .. note:: + If :attr:`MultiMeshRayCasterCfg.reference_meshes` is False, this flag has no effect. + """ + + merge_prim_meshes: bool = True + """Whether to merge the parsed meshes for a prim that contains multiple meshes. Defaults to True. + + This will create a new mesh that combines all meshes in the parsed prim. The raycast hits mesh IDs will then refer to the single + merged mesh. + """ + + track_mesh_transforms: bool = True + """Whether the mesh transformations should be tracked. Defaults to True. + + .. note:: + Not tracking the mesh transformations is recommended when the meshes are static to increase performance. + """ + + class_type: type = MultiMeshRayCaster + + mesh_prim_paths: list[str | RaycastTargetCfg] = MISSING + """The list of mesh primitive paths to ray cast against. + + If an entry is a string, it is internally converted to :class:`RaycastTargetCfg` with :attr:`~RaycastTargetCfg.is_global` enabled and attr:`~RaycastTargetCfg.track_mesh_transforms` disabled. These settings ensure backwards compatibility with the default raycaster. + """ + + update_mesh_ids: bool = False + """Whether to update the mesh ids of the ray hits in the :attr:`data` container.""" + + reference_meshes: bool = True + """Whether to reference duplicated meshes instead of loading each one separately into memory. + Defaults to True. + + When enabled, the raycaster parses all meshes in all environments, but reuses references + for duplicates instead of storing multiple copies. This reduces memory footprint. + """ diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_data.py b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_data.py new file mode 100644 index 00000000000..e9f88ead7ae --- /dev/null +++ b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_data.py @@ -0,0 +1,20 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + + +"""Configuration for the ray-cast sensor.""" +import torch + +from .ray_caster_data import RayCasterData + + +class MultiMeshRayCasterData(RayCasterData): + """Data container for the multi-mesh ray-cast sensor.""" + + ray_mesh_ids: torch.Tensor = None + """The mesh ids of the ray hits. + + Shape is (N, B, 1), where N is the number of sensors, B is the number of rays + in the scan pattern per sensor.""" diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster.py b/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster.py index 8e2f6541fe9..0cbdc8b0bf3 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster.py @@ -9,8 +9,9 @@ import re import torch from collections.abc import Sequence -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, ClassVar +import isaacsim.core.utils.stage as stage_utils import omni.log import omni.physics.tensors.impl.api as physx import warp as wp @@ -51,12 +52,25 @@ class RayCaster(SensorBase): cfg: RayCasterCfg """The configuration parameters.""" + # Class variables to share meshes and mesh_views across instances + meshes: ClassVar[dict[str, wp.Mesh]] = {} + """A dictionary to store warp meshes for raycasting, shared across all instances. + + The keys correspond to the prim path for the meshes, and values are the corresponding warp Mesh objects.""" + mesh_views: ClassVar[dict[str, XFormPrim | physx.ArticulationView | physx.RigidBodyView]] = {} + """A dictionary to store mesh views for raycasting, shared across all instances. + + The keys correspond to the prim path for the mesh views, and values are the corresponding view objects.""" + _instance_count: ClassVar[int] = 0 + """A counter to track the number of RayCaster instances, used to manage class variable lifecycle.""" + def __init__(self, cfg: RayCasterCfg): """Initializes the ray-caster object. Args: cfg: The configuration parameters. """ + RayCaster._instance_count += 1 # check if sensor path is valid # note: currently we do not handle environment indices if there is a regex pattern in the leaf # For example, if the prim path is "/World/Sensor_[1,2]". @@ -71,8 +85,6 @@ def __init__(self, cfg: RayCasterCfg): super().__init__(cfg) # Create empty variables for storing output data self._data = RayCasterData() - # the warp meshes used for raycasting. - self.meshes: dict[str, wp.Mesh] = {} def __str__(self) -> str: """Returns: A string containing information about the instance.""" @@ -80,7 +92,7 @@ def __str__(self) -> str: f"Ray-caster @ '{self.cfg.prim_path}': \n" f"\tview type : {self._view.__class__}\n" f"\tupdate period (s) : {self.cfg.update_period}\n" - f"\tnumber of meshes : {len(self.meshes)}\n" + f"\tnumber of meshes : {len(RayCaster.meshes)}\n" f"\tnumber of sensors : {self._view.count}\n" f"\tnumber of rays/sensor: {self.num_rays}\n" f"\ttotal number of rays : {self.num_rays * self._view.count}" @@ -131,28 +143,16 @@ def reset(self, env_ids: Sequence[int] | None = None): def _initialize_impl(self): super()._initialize_impl() # obtain global simulation view + self._physics_sim_view = SimulationManager.get_physics_sim_view() - # check if the prim at path is an articulated or rigid prim - # we do this since for physics-based view classes we can access their data directly - # otherwise we need to use the xform view class which is slower - found_supported_prim_class = False prim = sim_utils.find_first_matching_prim(self.cfg.prim_path) if prim is None: - raise RuntimeError(f"Failed to find a prim at path expression: {self.cfg.prim_path}") - # create view based on the type of prim - if prim.HasAPI(UsdPhysics.ArticulationRootAPI): - self._view = self._physics_sim_view.create_articulation_view(self.cfg.prim_path.replace(".*", "*")) - found_supported_prim_class = True - elif prim.HasAPI(UsdPhysics.RigidBodyAPI): - self._view = self._physics_sim_view.create_rigid_body_view(self.cfg.prim_path.replace(".*", "*")) - found_supported_prim_class = True - else: - self._view = XFormPrim(self.cfg.prim_path, reset_xform_properties=False) - found_supported_prim_class = True - omni.log.warn(f"The prim at path {prim.GetPath().pathString} is not a physics prim! Using XFormPrim.") - # check if prim view class is found - if not found_supported_prim_class: - raise RuntimeError(f"Failed to find a valid prim view class for the prim paths: {self.cfg.prim_path}") + available_prims = ",".join([str(p.GetPath()) for p in stage_utils.get_current_stage().Traverse()]) + raise RuntimeError( + f"Failed to find a prim at path expression: {self.cfg.prim_path}. Available prims: {available_prims}" + ) + + self._view, self._offset = self._get_trackable_prim_view(self.cfg.prim_path) # load the meshes by parsing the stage self._initialize_warp_meshes() @@ -168,6 +168,10 @@ def _initialize_warp_meshes(self): # read prims to ray-cast for mesh_prim_path in self.cfg.mesh_prim_paths: + # check if mesh already casted into warp mesh + if mesh_prim_path in RayCaster.meshes: + continue + # check if the prim is a plane - handle PhysX plane as a special case # if a plane exists then we need to create an infinite mesh that is a plane mesh_prim = sim_utils.get_first_matching_child_prim( @@ -201,10 +205,10 @@ def _initialize_warp_meshes(self): # print info omni.log.info(f"Created infinite plane mesh prim: {mesh_prim.GetPath()}.") # add the warp mesh to the list - self.meshes[mesh_prim_path] = wp_mesh + RayCaster.meshes[mesh_prim_path] = wp_mesh # throw an error if no meshes are found - if all([mesh_prim_path not in self.meshes for mesh_prim_path in self.cfg.mesh_prim_paths]): + if all([mesh_prim_path not in RayCaster.meshes for mesh_prim_path in self.cfg.mesh_prim_paths]): raise RuntimeError( f"No meshes found for ray-casting! Please check the mesh prim paths: {self.cfg.mesh_prim_paths}" ) @@ -225,26 +229,17 @@ def _initialize_rays_impl(self): self.drift = torch.zeros(self._view.count, 3, device=self.device) self.ray_cast_drift = torch.zeros(self._view.count, 3, device=self.device) # fill the data buffer - self._data.pos_w = torch.zeros(self._view.count, 3, device=self._device) - self._data.quat_w = torch.zeros(self._view.count, 4, device=self._device) - self._data.ray_hits_w = torch.zeros(self._view.count, self.num_rays, 3, device=self._device) + self._data.pos_w = torch.zeros(self._view.count, 3, device=self.device) + self._data.quat_w = torch.zeros(self._view.count, 4, device=self.device) + self._data.ray_hits_w = torch.zeros(self._view.count, self.num_rays, 3, device=self.device) + self._ray_starts_w = torch.zeros(self._view.count, self.num_rays, 3, device=self.device) + self._ray_directions_w = torch.zeros(self._view.count, self.num_rays, 3, device=self.device) - def _update_buffers_impl(self, env_ids: Sequence[int]): - """Fills the buffers of the sensor data.""" - # obtain the poses of the sensors - if isinstance(self._view, XFormPrim): - pos_w, quat_w = self._view.get_world_poses(env_ids) - elif isinstance(self._view, physx.ArticulationView): - pos_w, quat_w = self._view.get_root_transforms()[env_ids].split([3, 4], dim=-1) - quat_w = convert_quat(quat_w, to="wxyz") - elif isinstance(self._view, physx.RigidBodyView): - pos_w, quat_w = self._view.get_transforms()[env_ids].split([3, 4], dim=-1) - quat_w = convert_quat(quat_w, to="wxyz") - else: - raise RuntimeError(f"Unsupported view type: {type(self._view)}") - # note: we clone here because we are read-only operations - pos_w = pos_w.clone() - quat_w = quat_w.clone() + def _update_ray_infos(self, env_ids: Sequence[int]): + """Updates the ray information buffers.""" + + pos_w, quat_w = sim_utils.obtain_world_pose_from_view(self._view, env_ids) + pos_w, quat_w = math_utils.combine_frame_transforms(pos_w, quat_w, self._offset[0][env_ids], self._offset[1][env_ids]) # apply drift to ray starting position in world frame pos_w += self.drift[env_ids] # store the poses @@ -291,13 +286,20 @@ def _update_buffers_impl(self, env_ids: Sequence[int]): else: raise RuntimeError(f"Unsupported ray_alignment type: {self.cfg.ray_alignment}.") + self._ray_starts_w[env_ids] = ray_starts_w + self._ray_directions_w[env_ids] = ray_directions_w + + def _update_buffers_impl(self, env_ids: Sequence[int]): + """Fills the buffers of the sensor data.""" + self._update_ray_infos(env_ids) + # ray cast and store the hits # TODO: Make this work for multiple meshes? self._data.ray_hits_w[env_ids] = raycast_mesh( - ray_starts_w, - ray_directions_w, + self._ray_starts_w[env_ids], + self._ray_directions_w[env_ids], max_dist=self.cfg.max_distance, - mesh=self.meshes[self.cfg.mesh_prim_paths[0]], + mesh=RayCaster.meshes[self.cfg.mesh_prim_paths[0]], )[0] # apply vertical drift to ray starting position in ray caster frame @@ -316,12 +318,72 @@ def _set_debug_vis_impl(self, debug_vis: bool): self.ray_visualizer.set_visibility(False) def _debug_vis_callback(self, event): + if self._data.ray_hits_w is None: + return # remove possible inf values viz_points = self._data.ray_hits_w.reshape(-1, 3) viz_points = viz_points[~torch.any(torch.isinf(viz_points), dim=1)] - # show ray hit positions + self.ray_visualizer.visualize(viz_points) + def _get_trackable_prim_view( + self, target_prim_path: str + ) -> tuple[XFormPrim | any, tuple[torch.Tensor, torch.Tensor]]: + """Get a prim view that can be used to track the pose of the mesh prims. Additionally, it resolves the + relative pose between the mesh and its corresponding physics prim. This is especially useful if the + mesh is not directly parented to the physics prim. + """ + + mesh_prim = sim_utils.find_first_matching_prim(target_prim_path) + current_prim = mesh_prim + current_path_expr = target_prim_path + + prim_view = None + + while prim_view is None: + if current_prim.HasAPI(UsdPhysics.ArticulationRootAPI): + prim_view = self._physics_sim_view.create_articulation_view(current_path_expr.replace(".*", "*")) + omni.log.info(f"Created articulation view for mesh prim at path: {target_prim_path}") + break + + if current_prim.HasAPI(UsdPhysics.RigidBodyAPI): + prim_view = self._physics_sim_view.create_rigid_body_view(current_path_expr.replace(".*", "*")) + omni.log.info(f"Created rigid body view for mesh prim at path: {target_prim_path}") + break + + new_root_prim = current_prim.GetParent() + current_path_expr = current_path_expr.rsplit("/", 1)[0] + if not new_root_prim.IsValid(): + prim_view = XFormPrim(target_prim_path, reset_xform_properties=False) + current_path_expr = target_prim_path + omni.log.warn( + f"The prim at path {target_prim_path} which is used for raycasting is not a physics prim." + " Defaulting to XFormPrim. \n The pose of the mesh will most likely not" + " be updated correctly when running in headless mode and position lookups will be much slower. \n" + " If possible, ensure that the mesh or its parent is a physics prim (rigid body or articulation)." + ) + break + current_prim = new_root_prim + + mesh_prims = sim_utils.find_matching_prims(target_prim_path) + target_prims = sim_utils.find_matching_prims(current_path_expr) + if len(mesh_prims) != len(target_prims): + raise RuntimeError( + f"The number of mesh prims ({len(mesh_prims)}) does not match the number of physics prims" + f" ({len(target_prims)})Please specify the correct mesh and physics prim paths more" + " specifically in your target expressions." + ) + positions = [] + quaternions = [] + for mesh, target in zip(mesh_prims, target_prims): + pos, orientation = sim_utils.resolve_prim_pose(mesh, target) + positions.append(torch.tensor(pos, dtype=torch.float32, device=self.device)) + quaternions.append(torch.tensor(orientation, dtype=torch.float32, device=self.device)) + + positions = torch.stack(positions).to(device=self.device, dtype=torch.float32) + quaternions = torch.stack(quaternions).to(device=self.device, dtype=torch.float32) + return prim_view, (positions, quaternions) + """ Internal simulation callbacks. """ @@ -332,3 +394,9 @@ def _invalidate_initialize_callback(self, event): super()._invalidate_initialize_callback(event) # set all existing views to None to invalidate them self._view = None + + def __del__(self): + RayCaster._instance_count -= 1 + if RayCaster._instance_count == 0: + RayCaster.meshes.clear() + RayCaster.mesh_views.clear() diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_camera.py b/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_camera.py index 3bf8729b78b..07eb69e8408 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_camera.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_camera.py @@ -10,9 +10,9 @@ from typing import TYPE_CHECKING, ClassVar, Literal import isaacsim.core.utils.stage as stage_utils -import omni.physics.tensors.impl.api as physx -from isaacsim.core.prims import XFormPrim +import omni.log +import isaaclab.sim as sim_utils import isaaclab.utils.math as math_utils from isaaclab.sensors.camera import CameraData from isaaclab.utils.warp import raycast_mesh @@ -86,7 +86,7 @@ def __str__(self) -> str: f"Ray-Caster-Camera @ '{self.cfg.prim_path}': \n" f"\tview type : {self._view.__class__}\n" f"\tupdate period (s) : {self.cfg.update_period}\n" - f"\tnumber of meshes : {len(self.meshes)}\n" + f"\tnumber of meshes : {len(RayCaster.meshes)}\n" f"\tnumber of sensors : {self._view.count}\n" f"\tnumber of rays/sensor: {self.num_rays}\n" f"\ttotal number of rays : {self.num_rays * self._view.count}\n" @@ -143,11 +143,14 @@ def reset(self, env_ids: Sequence[int] | None = None): # reset the timestamps super().reset(env_ids) # resolve None - if env_ids is None: - env_ids = slice(None) + if env_ids is None or isinstance(env_ids, slice): + env_ids = self._ALL_INDICES # reset the data # note: this recomputation is useful if one performs events such as randomizations on the camera poses. - pos_w, quat_w = self._compute_camera_world_poses(env_ids) + pos_w, quat_w = sim_utils.obtain_world_pose_from_view(self._view, env_ids, clone=True) + pos_w, quat_w = math_utils.combine_frame_transforms( + pos_w, quat_w, self._offset_pos[env_ids], self._offset_quat[env_ids] + ) self._data.pos_w[env_ids] = pos_w self._data.quat_w_world[env_ids] = quat_w # Reset the frame count @@ -184,11 +187,11 @@ def set_world_poses( RuntimeError: If the camera prim is not set. Need to call :meth:`initialize` method first. """ # resolve env_ids - if env_ids is None: + if env_ids is None or isinstance(env_ids, slice): env_ids = self._ALL_INDICES # get current positions - pos_w, quat_w = self._compute_view_world_poses(env_ids) + pos_w, quat_w = sim_utils.obtain_world_pose_from_view(self._view, env_ids) if positions is not None: # transform to camera frame pos_offset_world_frame = positions - pos_w @@ -201,7 +204,10 @@ def set_world_poses( self._offset_quat[env_ids] = math_utils.quat_mul(math_utils.quat_inv(quat_w), quat_w_set) # update the data - pos_w, quat_w = self._compute_camera_world_poses(env_ids) + pos_w, quat_w = sim_utils.obtain_world_pose_from_view(self._view, env_ids, clone=True) + pos_w, quat_w = math_utils.combine_frame_transforms( + pos_w, quat_w, self._offset_pos[env_ids], self._offset_quat[env_ids] + ) self._data.pos_w[env_ids] = pos_w self._data.quat_w_world[env_ids] = quat_w @@ -260,7 +266,10 @@ def _update_buffers_impl(self, env_ids: Sequence[int]): self._frame[env_ids] += 1 # compute poses from current view - pos_w, quat_w = self._compute_camera_world_poses(env_ids) + pos_w, quat_w = sim_utils.obtain_world_pose_from_view(self._view, env_ids, clone=True) + pos_w, quat_w = math_utils.combine_frame_transforms( + pos_w, quat_w, self._offset_pos[env_ids], self._offset_quat[env_ids] + ) # update the data self._data.pos_w[env_ids] = pos_w self._data.quat_w_world[env_ids] = quat_w @@ -280,7 +289,7 @@ def _update_buffers_impl(self, env_ids: Sequence[int]): self.ray_hits_w, ray_depth, ray_normal, _ = raycast_mesh( ray_starts_w, ray_directions_w, - mesh=self.meshes[self.cfg.mesh_prim_paths[0]], + mesh=RayCaster.meshes[self.cfg.mesh_prim_paths[0]], max_dist=1e6, return_distance=any( [name in self.cfg.data_types for name in ["distance_to_image_plane", "distance_to_camera"]] @@ -395,39 +404,46 @@ def _compute_intrinsic_matrices(self): def _compute_view_world_poses(self, env_ids: Sequence[int]) -> tuple[torch.Tensor, torch.Tensor]: """Obtains the pose of the view the camera is attached to in the world frame. + .. deprecated v2.3.0: + This function will be removed in a future release in favor of implementation :meth:`sim_utils.obtain_world_pose_from_view`. + Returns: A tuple of the position (in meters) and quaternion (w, x, y, z). + + """ - # obtain the poses of the sensors - # note: clone arg doesn't exist for xform prim view so we need to do this manually - if isinstance(self._view, XFormPrim): - if isinstance(env_ids, slice): # catch the case where env_ids is a slice - env_ids = self._ALL_INDICES - pos_w, quat_w = self._view.get_world_poses(env_ids) - elif isinstance(self._view, physx.ArticulationView): - pos_w, quat_w = self._view.get_root_transforms()[env_ids].split([3, 4], dim=-1) - quat_w = math_utils.convert_quat(quat_w, to="wxyz") - elif isinstance(self._view, physx.RigidBodyView): - pos_w, quat_w = self._view.get_transforms()[env_ids].split([3, 4], dim=-1) - quat_w = math_utils.convert_quat(quat_w, to="wxyz") - else: - raise RuntimeError(f"Unsupported view type: {type(self._view)}") - # return the pose - return pos_w.clone(), quat_w.clone() + # deprecation + omni.log.warn( + "The function '_compute_view_world_poses' will be deprecated in favor of the util method" + " 'sim_utils.obtain_world_pose_from_view'. Please use 'sim_utils.obtain_world_pose_from_view' instead...." + ) + + return sim_utils.obtain_world_pose_from_view(self._view, env_ids, clone=True) def _compute_camera_world_poses(self, env_ids: Sequence[int]) -> tuple[torch.Tensor, torch.Tensor]: """Computes the pose of the camera in the world frame. This function applies the offset pose to the pose of the view the camera is attached to. + .. deprecated v2.3.0: + This function will be removed in a future release. Instead, use the code block below: + + .. code-block:: python + + pos_w, quat_w = sim_utils.obtain_world_pose_from_view(self._view, env_ids, clone=True) + pos_w, quat_w = math_utils.combine_frame_transforms(pos_w, quat_w, self._offset_pos[env_ids], self._offset_quat[env_ids]) + Returns: A tuple of the position (in meters) and quaternion (w, x, y, z) in "world" convention. """ - # get the pose of the view the camera is attached to - pos_w, quat_w = self._compute_view_world_poses(env_ids) - # apply offsets - # need to apply quat because offset relative to parent frame - pos_w += math_utils.quat_apply(quat_w, self._offset_pos[env_ids]) - quat_w = math_utils.quat_mul(quat_w, self._offset_quat[env_ids]) - return pos_w, quat_w + # deprecation + omni.log.warn( + "The function '_compute_camera_world_poses' will be deprecated in favor of the combination of methods" + " 'sim_utils.obtain_world_pose_from_view' and 'math_utils.combine_frame_transforms'. Please use" + " 'sim_utils.obtain_world_pose_from_view' and 'math_utils.combine_frame_transforms' instead...." + ) + + # get the pose of the view the camera is attached to + pos_w, quat_w = sim_utils.obtain_world_pose_from_view(self._view, env_ids, clone=True) + return math_utils.combine_frame_transforms(pos_w, quat_w, self._offset_pos[env_ids], self._offset_quat[env_ids]) diff --git a/source/isaaclab/isaaclab/sim/utils.py b/source/isaaclab/isaaclab/sim/utils.py index 3cdbc182a6f..8f47dc70138 100644 --- a/source/isaaclab/isaaclab/sim/utils.py +++ b/source/isaaclab/isaaclab/sim/utils.py @@ -11,6 +11,7 @@ import functools import inspect import re +import torch from collections.abc import Callable, Generator from typing import TYPE_CHECKING, Any @@ -19,12 +20,16 @@ import omni import omni.kit.commands import omni.log +import omni.physics.tensors.impl.api as physx from isaacsim.core.cloner import Cloner +from isaacsim.core.prims import XFormPrim from isaacsim.core.utils.carb import get_carb_setting from isaacsim.core.utils.stage import get_current_stage from isaacsim.core.version import get_version from pxr import PhysxSchema, Sdf, Usd, UsdGeom, UsdPhysics, UsdShade, UsdUtils +from isaaclab.utils.math import convert_quat + # from Isaac Sim 4.2 onwards, pxr.Semantics is deprecated try: import Semantics @@ -571,6 +576,92 @@ def make_uninstanceable(prim_path: str | Sdf.Path, stage: Usd.Stage | None = Non all_prims += child_prim.GetFilteredChildren(Usd.TraverseInstanceProxies()) +def resolve_prim_pose( + prim: Usd.Prim, ref_prim: Usd.Prim | None = None +) -> tuple[tuple[float, float, float], tuple[float, float, float, float]]: + """Resolve the pose of a prim with respect to another prim. + + Note: + This function ignores scale and skew by orthonormalizing the transformation + matrix at the final step. However, if any ancestor prim in the hierarchy + has non-uniform scale, that scale will still affect the resulting position + and orientation of the prim (because it's baked into the transform before + scale removal). + + In other words: scale **is not removed hierarchically**. If you need + completely scale-free poses, you must walk the transform chain and strip + scale at each level. Please open an issue if you need this functionality. + + Args: + prim: The USD prim to resolve the pose for. + ref_prim: The USD prim to compute the pose with respect to. + Defaults to None, in which case the world frame is used. + + Returns: + A tuple containing the position (as a 3D vector) and the quaternion orientation + in the (w, x, y, z) format. + + Raises: + ValueError: If the prim or ref prim is not valid. + """ + # check if prim is valid + if not prim.IsValid(): + raise ValueError(f"Prim at path '{prim.GetPath().pathString}' is not valid.") + # get prim xform + xform = UsdGeom.Xformable(prim) + prim_tf = xform.ComputeLocalToWorldTransform(Usd.TimeCode.Default()) + # sanitize quaternion + # this is needed, otherwise the quaternion might be non-normalized + prim_tf = prim_tf.GetOrthonormalized() + + if ref_prim is not None: + # check if ref prim is valid + if not ref_prim.IsValid(): + raise ValueError(f"Ref prim at path '{ref_prim.GetPath().pathString}' is not valid.") + # get ref prim xform + ref_xform = UsdGeom.Xformable(ref_prim) + ref_tf = ref_xform.ComputeLocalToWorldTransform(Usd.TimeCode.Default()) + # make sure ref tf is orthonormal + ref_tf = ref_tf.GetOrthonormalized() + # compute relative transform to get prim in ref frame + prim_tf = prim_tf * ref_tf.GetInverse() + + # extract position and orientation + prim_pos = [*prim_tf.ExtractTranslation()] + prim_quat = [prim_tf.ExtractRotationQuat().real, *prim_tf.ExtractRotationQuat().imaginary] + return tuple(prim_pos), tuple(prim_quat) + + +def resolve_prim_scale(prim: Usd.Prim) -> tuple[float, float, float]: + """Resolve the scale of a prim in the world frame. + + At an attribute level, a USD prim's scale is a scaling transformation applied to the prim with + respect to its parent prim. This function resolves the scale of the prim in the world frame, + by computing the local to world transform of the prim. This is equivalent to traversing up + the prim hierarchy and accounting for the rotations and scales of the prims. + + For instance, if a prim has a scale of (1, 2, 3) and it is a child of a prim with a scale of (4, 5, 6), + then the scale of the prim in the world frame is (4, 10, 18). + + Args: + prim: The USD prim to resolve the scale for. + + Returns: + The scale of the prim in the x, y, and z directions in the world frame. + + Raises: + ValueError: If the prim is not valid. + """ + # check if prim is valid + if not prim.IsValid(): + raise ValueError(f"Prim at path '{prim.GetPath().pathString}' is not valid.") + # compute local to world transform + xform = UsdGeom.Xformable(prim) + world_transform = xform.ComputeLocalToWorldTransform(Usd.TimeCode.Default()) + # extract scale + return tuple([*(v.GetLength() for v in world_transform.ExtractRotationMatrix())]) + + """ USD Stage traversal. """ @@ -1093,3 +1184,42 @@ def get_current_stage_id() -> int: if stage_id < 0: stage_id = stage_cache.Insert(stage).ToLongInt() return stage_id + + +""" +PhysX prim views utils. +""" + + +def obtain_world_pose_from_view( + physx_view: XFormPrim | physx.ArticulationView | physx.RigidBodyView, + env_ids: torch.Tensor, + clone: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + """Get the world poses of the prim referenced by the prim view. + + Args: + physx_view: The prim view to get the world poses from. + env_ids: The environment ids of the prims to get the world poses for. + + Returns: + A tuple containing the world positions and orientations of the prims. Orientation is in wxyz format. + + Raises: + NotImplementedError: If the prim view is not of the correct type. + """ + if isinstance(physx_view, XFormPrim): + pos_w, quat_w = physx_view.get_world_poses(env_ids) + elif isinstance(physx_view, physx.ArticulationView): + pos_w, quat_w = physx_view.get_root_transforms()[env_ids].split([3, 4], dim=-1) + quat_w = convert_quat(quat_w, to="wxyz") + elif isinstance(physx_view, physx.RigidBodyView): + pos_w, quat_w = physx_view.get_transforms()[env_ids].split([3, 4], dim=-1) + quat_w = convert_quat(quat_w, to="wxyz") + else: + raise NotImplementedError(f"Cannot get world poses for prim view of type '{type(physx_view)}'.") + + if clone: + return pos_w.clone(), quat_w.clone() + else: + return pos_w, quat_w diff --git a/source/isaaclab/isaaclab/utils/__init__.py b/source/isaaclab/isaaclab/utils/__init__.py index a5365948699..036d7c9c0c6 100644 --- a/source/isaaclab/isaaclab/utils/__init__.py +++ b/source/isaaclab/isaaclab/utils/__init__.py @@ -10,6 +10,7 @@ from .configclass import configclass from .dict import * from .interpolation import * +from .mesh import * from .modifiers import * from .string import * from .timer import Timer diff --git a/source/isaaclab/isaaclab/utils/mesh.py b/source/isaaclab/isaaclab/utils/mesh.py new file mode 100644 index 00000000000..4f69b8cd9cc --- /dev/null +++ b/source/isaaclab/isaaclab/utils/mesh.py @@ -0,0 +1,132 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + + +"""Utility functions for working with meshes.""" + +import numpy as np +import trimesh + +from pxr import Usd, UsdGeom + + +def create_trimesh_from_geom_mesh(mesh_prim: Usd.Prim) -> trimesh.Trimesh: + """Reads the vertices and faces of a mesh prim. + + The function reads the vertices and faces of a mesh prim and returns it. If the underlying mesh is a quad mesh, + it converts it to a triangle mesh. + + Args: + mesh_prim: The mesh prim to read the vertices and faces from. + + Returns: + A tuple containing the vertices and faces of the mesh. + Shape of vertices is (n_vertices, 3). + Shape of faces is (n_faces, 3). + """ + if mesh_prim.GetTypeName() != "Mesh": + raise ValueError(f"Prim at path '{mesh_prim.GetPath()}' is not a mesh.") + # cast into UsdGeomMesh + mesh = UsdGeom.Mesh(mesh_prim) + + # read the vertices and faces + points = np.asarray(mesh.GetPointsAttr().Get()).copy() + + # Load faces and convert to triangle if needed. (Default is quads) + num_vertex_per_face = np.asarray(mesh.GetFaceVertexCountsAttr().Get()) + indices = np.asarray(mesh.GetFaceVertexIndicesAttr().Get()) + return trimesh.Trimesh(points, convert_faces_to_triangles(indices, num_vertex_per_face)) + + +def create_trimesh_from_geom_shape(prim: Usd.Prim) -> trimesh.Trimesh: + """Converts a primitive object to a trimesh. + + Args: + prim: The prim that should be converted to a trimesh. + + Returns: + A trimesh object representing the primitive. + + Raises: + ValueError: If the prim is not a supported primitive. Check PRIMITIVE_MESH_TYPES for supported primitives. + """ + + if prim.GetTypeName() not in PRIMITIVE_MESH_TYPES: + raise ValueError(f"Prim at path '{prim.GetPath()}' is not a primitive mesh. Cannot convert to trimesh.") + + return _MESH_CONVERTERS_CALLBACKS[prim.GetTypeName()](prim) + + +def convert_faces_to_triangles(faces: np.ndarray, point_counts: np.ndarray) -> np.ndarray: + """Converts quad mesh face indices into triangle face indices. + + This function expects an array of faces (indices) and the number of points per face. It then converts potential + quads into triangles and returns the new triangle face indices as a numpy array of shape (n_faces_new, 3). + + Args: + faces: The faces of the quad mesh as a one-dimensional array. Shape is (N,). + point_counts: The number of points per face. Shape is (N, 3) or (N, 4). + + Returns: + The new face ids with triangles. Shape is (n_faces_new, 3). + """ + # check if the mesh is already triangulated + if (point_counts == 3).all(): + return faces.reshape(-1, 3) # already triangulated + all_faces = [] + + vertex_counter = 0 + # Iterates over all triangles of the mesh. + # could be very slow for large meshes + for num_points in point_counts: + if num_points == 3: + # triangle + all_faces.append(faces[vertex_counter : vertex_counter + 3]) + elif num_points == 4: + # quads. Subdivide into two triangles + f = faces[vertex_counter : vertex_counter + 4] + first_triangle = f[:3] + second_triangle = np.array([f[0], f[2], f[3]]) + all_faces.append(first_triangle) + all_faces.append(second_triangle) + else: + raise RuntimeError(f"Invalid number of points per face: {num_points}") + + vertex_counter += num_points + return np.asarray(all_faces) + + +def _create_plane_trimesh(prim: Usd.Prim) -> trimesh.Trimesh: + """Creates a trimesh for a plane primitive.""" + size = (2e6, 2e6) + vertices = np.array([[size[0], size[1], 0], [size[0], 0.0, 0], [0.0, size[1], 0], [0.0, 0.0, 0]]) - np.array( + [size[0] / 2.0, size[1] / 2.0, 0.0] + ) + faces = np.array([[1, 0, 2], [2, 3, 1]]) + return trimesh.Trimesh(vertices=vertices, faces=faces) + + +def _create_cube_trimesh(prim: Usd.Prim) -> trimesh.Trimesh: + """Creates a trimesh for a cube primitive.""" + size = prim.GetAttribute("size").Get() + extends = [size, size, size] + return trimesh.creation.box(extends) + + +def _create_sphere_trimesh(prim: Usd.Prim, subdivisions: int = 2) -> trimesh.Trimesh: + """Creates a trimesh for a sphere primitive.""" + radius = prim.GetAttribute("radius").Get() + mesh = trimesh.creation.icosphere(radius=radius, subdivisions=subdivisions) + return mesh + + +_MESH_CONVERTERS_CALLBACKS: dict[str, callable] = { + "Plane": _create_plane_trimesh, + "Cube": _create_cube_trimesh, + "Sphere": _create_sphere_trimesh, +} + +PRIMITIVE_MESH_TYPES = list(_MESH_CONVERTERS_CALLBACKS.keys()) +"""List of supported primitive mesh types that can be converted to a trimesh.""" diff --git a/source/isaaclab/isaaclab/utils/warp/__init__.py b/source/isaaclab/isaaclab/utils/warp/__init__.py index 14c49f25528..8400fb670a0 100644 --- a/source/isaaclab/isaaclab/utils/warp/__init__.py +++ b/source/isaaclab/isaaclab/utils/warp/__init__.py @@ -5,4 +5,4 @@ """Sub-module containing operations based on warp.""" -from .ops import convert_to_warp_mesh, raycast_mesh +from .ops import convert_to_warp_mesh, raycast_dynamic_meshes, raycast_mesh, raycast_single_mesh diff --git a/source/isaaclab/isaaclab/utils/warp/kernels.py b/source/isaaclab/isaaclab/utils/warp/kernels.py index 2fe544651f5..23fb2c5836e 100644 --- a/source/isaaclab/isaaclab/utils/warp/kernels.py +++ b/source/isaaclab/isaaclab/utils/warp/kernels.py @@ -75,6 +75,173 @@ def raycast_mesh_kernel( ray_face_id[tid] = f +@wp.kernel(enable_backward=False) +def raycast_static_meshes_kernel( + mesh: wp.array2d(dtype=wp.uint64), + ray_starts: wp.array2d(dtype=wp.vec3), + ray_directions: wp.array2d(dtype=wp.vec3), + ray_hits: wp.array2d(dtype=wp.vec3), + ray_distance: wp.array2d(dtype=wp.float32), + ray_normal: wp.array2d(dtype=wp.vec3), + ray_face_id: wp.array2d(dtype=wp.int32), + ray_mesh_id: wp.array2d(dtype=wp.int16), + max_dist: float = 1e6, + return_normal: int = False, + return_face_id: int = False, + return_mesh_id: int = False, +): + """Performs ray-casting against multiple static meshes. + + This function performs ray-casting against the given meshes using the provided ray start positions + and directions. The resulting ray hit positions are stored in the :obj:`ray_hits` array. + + The function utilizes the ``mesh_query_ray`` method from the ``wp`` module to perform the actual ray-casting + operation. The maximum ray-cast distance is set to ``1e6`` units. + + .. note:: + That the ``ray_starts``, ``ray_directions``, and ``ray_hits`` arrays should have compatible shapes + and data types to ensure proper execution. Additionally, they all must be in the same frame. + + This kernel differs from the :meth:`raycast_dynamic_meshes_kernel` in that it does not take into + account the mesh's position and rotation. This kernel is useful for ray-casting against static meshes + that are not expected to move. + + Args: + mesh: The input mesh. The ray-casting is performed against this mesh on the device specified by the + `mesh`'s `device` attribute. + ray_starts: The input ray start positions. Shape is (B, N, 3). + ray_directions: The input ray directions. Shape is (B, N, 3). + ray_hits: The output ray hit positions. Shape is (B, N, 3). + ray_distance: The output ray hit distances. Shape is (B, N,), if ``return_distance`` is True. Otherwise, + this array is not used. + ray_normal: The output ray hit normals. Shape is (B, N, 3), if ``return_normal`` is True. Otherwise, + this array is not used. + ray_face_id: The output ray hit face ids. Shape is (B, N,), if ``return_face_id`` is True. Otherwise, + this array is not used. + ray_mesh_id: The output ray hit mesh ids. Shape is (B, N,), if ``return_mesh_id`` is True. Otherwise, + this array is not used. + max_dist: The maximum ray-cast distance. Defaults to 1e6. + return_normal: Whether to return the ray hit normals. Defaults to False`. + return_face_id: Whether to return the ray hit face ids. Defaults to False. + return_mesh_id: Whether to return the mesh id. Defaults to False. + """ + # get the thread id + tid_mesh_id, tid_env, tid_ray = wp.tid() + + direction = ray_directions[tid_env, tid_ray] + start_pos = ray_starts[tid_env, tid_ray] + + # ray cast against the mesh and store the hit position + mesh_query_ray_t = wp.mesh_query_ray(mesh[tid_env, tid_mesh_id], start_pos, direction, max_dist) + + # if the ray hit, store the hit data + if mesh_query_ray_t.result: + + # check if hit distance is less than the current hit distance, only then update the memory + if mesh_query_ray_t.t < ray_distance[tid_env, tid_ray]: + + # convert back to world space and update the hit data + ray_hits[tid_env, tid_ray] = start_pos + mesh_query_ray_t.t * direction + + # update the hit distance + ray_distance[tid_env, tid_ray] = mesh_query_ray_t.t + + # update the normal and face id if requested + if return_normal == 1: + ray_normal[tid_env, tid_ray] = mesh_query_ray_t.normal + if return_face_id == 1: + ray_face_id[tid_env, tid_ray] = mesh_query_ray_t.face + if return_mesh_id == 1: + ray_mesh_id[tid_env, tid_ray] = wp.int16(tid_mesh_id) + + +@wp.kernel +def raycast_dynamic_meshes_kernel( + mesh: wp.array2d(dtype=wp.uint64), + ray_starts: wp.array2d(dtype=wp.vec3), + ray_directions: wp.array2d(dtype=wp.vec3), + ray_hits: wp.array2d(dtype=wp.vec3), + ray_distance: wp.array2d(dtype=wp.float32), + ray_normal: wp.array2d(dtype=wp.vec3), + ray_face_id: wp.array2d(dtype=wp.int32), + ray_mesh_id: wp.array2d(dtype=wp.int16), + mesh_positions: wp.array2d(dtype=wp.vec3), + mesh_rotations: wp.array2d(dtype=wp.quat), + max_dist: float = 1e6, + return_normal: int = False, + return_face_id: int = False, + return_mesh_id: int = False, +): + """Performs ray-casting against multiple meshes. + + This function performs ray-casting against the given meshes using the provided ray start positions + and directions. The resulting ray hit positions are stored in the :obj:`ray_hits` array. + + The function utilizes the ``mesh_query_ray`` method from the ``wp`` module to perform the actual ray-casting + operation. The maximum ray-cast distance is set to ``1e6`` units. + + + Note: + That the ``ray_starts``, ``ray_directions``, and ``ray_hits`` arrays should have compatible shapes + and data types to ensure proper execution. Additionally, they all must be in the same frame. + + All arguments are expected to be batched with the first dimension (B, batch) being the number of envs + and the second dimension (N, num_rays) being the number of rays. For Meshes, W is the number of meshes. + + Args: + mesh: The input mesh. The ray-casting is performed against this mesh on the device specified by the + `mesh`'s `device` attribute. + ray_starts: The input ray start positions. Shape is (B, N, 3). + ray_directions: The input ray directions. Shape is (B, N, 3). + ray_hits: The output ray hit positions. Shape is (B, N, 3). + ray_distance: The output ray hit distances. Shape is (B, N,), if ``return_distance`` is True. Otherwise, + this array is not used. + ray_normal: The output ray hit normals. Shape is (B, N, 3), if ``return_normal`` is True. Otherwise, + this array is not used. + ray_face_id: The output ray hit face ids. Shape is (B, N,), if ``return_face_id`` is True. Otherwise, + this array is not used. + ray_mesh_id: The output ray hit mesh ids. Shape is (B, N,), if ``return_mesh_id`` is True. Otherwise, + this array is not used. + mesh_positions: The input mesh positions in world frame. Shape is (W, 3). + mesh_rotations: The input mesh rotations in world frame. Shape is (W, 3, 3). + max_dist: The maximum ray-cast distance. Defaults to 1e6. + return_normal: Whether to return the ray hit normals. Defaults to False`. + return_face_id: Whether to return the ray hit face ids. Defaults to False. + return_mesh_id: Whether to return the mesh id. Defaults to False. + """ + # get the thread id + tid_mesh_id, tid_env, tid_ray = wp.tid() + + mesh_pose = wp.transform(mesh_positions[tid_env, tid_mesh_id], mesh_rotations[tid_env, tid_mesh_id]) + mesh_pose_inv = wp.transform_inverse(mesh_pose) + direction = wp.transform_vector(mesh_pose_inv, ray_directions[tid_env, tid_ray]) + start_pos = wp.transform_point(mesh_pose_inv, ray_starts[tid_env, tid_ray]) + + # ray cast against the mesh and store the hit position + mesh_query_ray_t = wp.mesh_query_ray(mesh[tid_env, tid_mesh_id], start_pos, direction, max_dist) + # if the ray hit, store the hit data + if mesh_query_ray_t.result: + + # check if hit distance is less than the current hit distance, only then update the memory + if mesh_query_ray_t.t < ray_distance[tid_env, tid_ray]: + + # convert back to world space and update the hit data + hit_pos = start_pos + mesh_query_ray_t.t * direction + ray_hits[tid_env, tid_ray] = wp.transform_point(mesh_pose, hit_pos) + + # update the hit distance + ray_distance[tid_env, tid_ray] = mesh_query_ray_t.t + + # update the normal and face id if requested + if return_normal == 1: + n = wp.transform_vector(mesh_pose, mesh_query_ray_t.normal) + ray_normal[tid_env, tid_ray] = n + if return_face_id == 1: + ray_face_id[tid_env, tid_ray] = mesh_query_ray_t.face + if return_mesh_id == 1: + ray_mesh_id[tid_env, tid_ray] = wp.int16(tid_mesh_id) + + @wp.kernel(enable_backward=False) def reshape_tiled_image( tiled_image_buffer: Any, diff --git a/source/isaaclab/isaaclab/utils/warp/ops.py b/source/isaaclab/isaaclab/utils/warp/ops.py index a2db46c4b52..41f363a3997 100644 --- a/source/isaaclab/isaaclab/utils/warp/ops.py +++ b/source/isaaclab/isaaclab/utils/warp/ops.py @@ -18,6 +18,8 @@ # initialize the warp module wp.init() +from isaaclab.utils.math import convert_quat + from . import kernels @@ -127,6 +129,257 @@ def raycast_mesh( return ray_hits.to(device).view(shape), ray_distance, ray_normal, ray_face_id +def raycast_single_mesh( + ray_starts: torch.Tensor, + ray_directions: torch.Tensor, + mesh_id: int, + max_dist: float = 1e6, + return_distance: bool = False, + return_normal: bool = False, + return_face_id: bool = False, +) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: + """Performs ray-casting against a mesh. + + Note that the :attr:`ray_starts` and :attr:`ray_directions`, and :attr:`ray_hits` should have compatible shapes + and data types to ensure proper execution. Additionally, they all must be in the same frame. + + Args: + ray_starts: The starting position of the rays. Shape (B, N, 3). + ray_directions: The ray directions for each ray. Shape (B, N, 3). + mesh_id: The warp mesh id to ray-cast against. + max_dist: The maximum distance to ray-cast. Defaults to 1e6. + return_distance: Whether to return the distance of the ray until it hits the mesh. Defaults to False. + return_normal: Whether to return the normal of the mesh face the ray hits. Defaults to False. + return_face_id: Whether to return the face id of the mesh face the ray hits. Defaults to False. + + Returns: + The ray hit position. Shape (B, N, 3). + The returned tensor contains :obj:`float('inf')` for missed hits. + The ray hit distance. Shape (B, N,). + Will only return if :attr:`return_distance` is True, else returns None. + The returned tensor contains :obj:`float('inf')` for missed hits. + The ray hit normal. Shape (B, N, 3). + Will only return if :attr:`return_normal` is True else returns None. + The returned tensor contains :obj:`float('inf')` for missed hits. + The ray hit face id. Shape (B, N,). + Will only return if :attr:`return_face_id` is True else returns None. + The returned tensor contains :obj:`int(-1)` for missed hits. + """ + # cast mesh id into array + mesh_ids = wp.array2d( + [[mesh_id] * ray_starts.shape[0]], dtype=wp.uint64, device=wp.device_from_torch(ray_starts.device) + ) + ray_hits, ray_distance, ray_normal, ray_face_id, ray_mesh_id = raycast_dynamic_meshes( + ray_starts=ray_starts, + ray_directions=ray_directions, + mesh_ids_wp=mesh_ids, + max_dist=max_dist, + return_distance=return_distance, + return_normal=return_normal, + return_face_id=return_face_id, + ) + + return ray_hits, ray_distance, ray_normal, ray_face_id + + +def raycast_dynamic_meshes( + ray_starts: torch.Tensor, + ray_directions: torch.Tensor, + mesh_ids_wp: wp.Array, + mesh_positions_w: torch.Tensor | None = None, + mesh_orientations_w: torch.Tensor | None = None, + max_dist: float = 1e6, + return_distance: bool = False, + return_normal: bool = False, + return_face_id: bool = False, + return_mesh_id: bool = False, +) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: + """Performs ray-casting against multiple, dynamic meshes. + + Note that the :attr:`ray_starts` and :attr:`ray_directions`, and :attr:`ray_hits` should have compatible shapes + and data types to ensure proper execution. Additionally, they all must be in the same frame. + + If mesh positions and rotations are provided, they need to have to have the same shape as the + number of meshes. + + Args: + ray_starts: The starting position of the rays. Shape (B, N, 3). + ray_directions: The ray directions for each ray. Shape (B, N, 3). + mesh_ids_wp: The warp mesh ids to ray-cast against. Length (B, M). + mesh_positions_w: The world positions of the meshes. Shape (B, M, 3). + mesh_orientations_w: The world orientation as quaternion (wxyz) format. Shape (B, M, 4). + max_dist: The maximum distance to ray-cast. Defaults to 1e6. + return_distance: Whether to return the distance of the ray until it hits the mesh. Defaults to False. + return_normal: Whether to return the normal of the mesh face the ray hits. Defaults to False. + return_face_id: Whether to return the face id of the mesh face the ray hits. Defaults to False. + return_mesh_id: Whether to return the mesh id of the mesh face the ray hits. Defaults to False. + NOTE: the type of the returned tensor is torch.int16, so you can't have more than 32767 meshes. + + Returns: + The ray hit position. Shape (B, N, 3). + The returned tensor contains :obj:`float('inf')` for missed hits. + The ray hit distance. Shape (B, N,). + Will only return if :attr:`return_distance` is True, else returns None. + The returned tensor contains :obj:`float('inf')` for missed hits. + The ray hit normal. Shape (B, N, 3). + Will only return if :attr:`return_normal` is True else returns None. + The returned tensor contains :obj:`float('inf')` for missed hits. + The ray hit face id. Shape (B, N,). + Will only return if :attr:`return_face_id` is True else returns None. + The returned tensor contains :obj:`int(-1)` for missed hits. + The ray hit mesh id. Shape (B, N,). + Will only return if :attr:`return_mesh_id` is True else returns None. + The returned tensor contains :obj:`0` for missed hits. + """ + # extract device and shape information + shape = ray_starts.shape + device = ray_starts.device + + # device of the mesh + torch_device = wp.device_to_torch(mesh_ids_wp.device) + n_meshes = mesh_ids_wp.shape[1] + + n_envs = ray_starts.shape[0] + n_rays_per_env = ray_starts.shape[1] + + # reshape the tensors + ray_starts = ray_starts.to(torch_device).view(n_envs, n_rays_per_env, 3).contiguous() + ray_directions = ray_directions.to(torch_device).view(n_envs, n_rays_per_env, 3).contiguous() + + # create output tensor for the ray hits + ray_hits = torch.full((n_envs, n_rays_per_env, 3), float("inf"), device=torch_device).contiguous() + + # map the memory to warp arrays + ray_starts_wp = wp.from_torch(ray_starts, dtype=wp.vec3) + ray_directions_wp = wp.from_torch(ray_directions, dtype=wp.vec3) + ray_hits_wp = wp.from_torch(ray_hits, dtype=wp.vec3) + # required to check if a closer hit is reported, returned only if return_distance is true + ray_distance = torch.full( + ( + n_envs, + n_rays_per_env, + ), + float("inf"), + device=torch_device, + ).contiguous() + ray_distance_wp = wp.from_torch(ray_distance, dtype=wp.float32) + + if return_normal: + ray_normal = torch.full((n_envs, n_rays_per_env, 3), float("inf"), device=torch_device).contiguous() + ray_normal_wp = wp.from_torch(ray_normal, dtype=wp.vec3) + else: + ray_normal = None + ray_normal_wp = wp.empty((1, 1), dtype=wp.vec3, device=torch_device) + + if return_face_id: + ray_face_id = torch.ones( + ( + n_envs, + n_rays_per_env, + ), + dtype=torch.int32, + device=torch_device, + ).contiguous() * (-1) + ray_face_id_wp = wp.from_torch(ray_face_id, dtype=wp.int32) + else: + ray_face_id = None + ray_face_id_wp = wp.empty((1, 1), dtype=wp.int32, device=torch_device) + + if return_mesh_id: + ray_mesh_id = -torch.ones((n_envs, n_rays_per_env), dtype=torch.int16, device=torch_device).contiguous() + ray_mesh_id_wp = wp.from_torch(ray_mesh_id, dtype=wp.int16) + else: + ray_mesh_id = None + ray_mesh_id_wp = wp.empty((1, 1), dtype=wp.int16, device=torch_device) + + ## + # Call the warp kernels + ### + if mesh_positions_w is None and mesh_orientations_w is None: + # Static mesh case, no need to pass in positions and rotations. + # launch the warp kernel + wp.launch( + kernel=kernels.raycast_static_meshes_kernel, + dim=[n_meshes, n_envs, n_rays_per_env], + inputs=[ + mesh_ids_wp, + ray_starts_wp, + ray_directions_wp, + ray_hits_wp, + ray_distance_wp, + ray_normal_wp, + ray_face_id_wp, + ray_mesh_id_wp, + float(max_dist), + int(return_normal), + int(return_face_id), + int(return_mesh_id), + ], + device=torch_device, + ) + else: + # dynamic mesh case + if mesh_positions_w is None: + mesh_positions_wp_w = wp.zeros((n_envs, n_meshes), dtype=wp.vec3, device=torch_device) + else: + mesh_positions_w = mesh_positions_w.to(torch_device).view(n_envs, n_meshes, 3).contiguous() + mesh_positions_wp_w = wp.from_torch(mesh_positions_w, dtype=wp.vec3) + + if mesh_orientations_w is None: + # Note (zrene): This is a little bit ugly, since it requires to initialize torch memory first + # But I couldn't find a better way to initialize a quaternion identity in warp + # wp.zeros(1, dtype=wp.quat, device=torch_device) gives all zero quaternion + quat_identity = torch.tensor([0, 0, 0, 1], dtype=torch.float32, device=torch_device).repeat( + n_envs, n_meshes, 1 + ) + mesh_quat_wp_w = wp.from_torch(quat_identity, dtype=wp.quat) + else: + mesh_orientations_w = convert_quat( + mesh_orientations_w.to(dtype=torch.float32, device=torch_device), "xyzw" + ).contiguous() + mesh_quat_wp_w = wp.from_torch(mesh_orientations_w, dtype=wp.quat) + + # launch the warp kernel + wp.launch( + kernel=kernels.raycast_dynamic_meshes_kernel, + dim=[n_meshes, n_envs, n_rays_per_env], + inputs=[ + mesh_ids_wp, + ray_starts_wp, + ray_directions_wp, + ray_hits_wp, + ray_distance_wp, + ray_normal_wp, + ray_face_id_wp, + ray_mesh_id_wp, + mesh_positions_wp_w, + mesh_quat_wp_w, + float(max_dist), + int(return_normal), + int(return_face_id), + int(return_mesh_id), + ], + device=torch_device, + ) + ## + # Cleanup and convert back to torch tensors + ## + + # NOTE: Synchronize is not needed anymore, but we keep it for now. Check with @dhoeller. + wp.synchronize() + + if return_distance: + ray_distance = ray_distance.to(device).view(shape[:2]) + if return_normal: + ray_normal = ray_normal.to(device).view(shape) + if return_face_id: + ray_face_id = ray_face_id.to(device).view(shape[:2]) + if return_mesh_id: + ray_mesh_id = ray_mesh_id.to(device).view(shape[:2]) + + return ray_hits.to(device).view(shape), ray_distance, ray_normal, ray_face_id, ray_mesh_id + + def convert_to_warp_mesh(points: np.ndarray, indices: np.ndarray, device: str) -> wp.Mesh: """Create a warp mesh object with a mesh defined from vertices and triangles. diff --git a/source/isaaclab/test/sensors/check_multi_mesh_ray_caster.py b/source/isaaclab/test/sensors/check_multi_mesh_ray_caster.py new file mode 100644 index 00000000000..f42f6a81ebd --- /dev/null +++ b/source/isaaclab/test/sensors/check_multi_mesh_ray_caster.py @@ -0,0 +1,214 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + + +""" +This script shows how to use the ray caster from the Isaac Lab framework. + +.. code-block:: bash + + # Usage + ./isaaclab.sh -p source/extensions/omni.isaac.lab/test/sensors/test_ray_caster.py --headless +""" + +"""Launch Isaac Sim Simulator first.""" + +import argparse + +from isaaclab.app import AppLauncher + +# add argparse arguments +parser = argparse.ArgumentParser(description="Ray Caster Test Script") +parser.add_argument("--num_envs", type=int, default=16, help="Number of environments to clone.") +parser.add_argument("--num_objects", type=int, default=0, help="Number of additional objects to clone.") +parser.add_argument( + "--terrain_type", + type=str, + default="generator", + help="Type of terrain to import. Can be 'generator' or 'usd' or 'plane'.", +) +# append AppLauncher cli args +AppLauncher.add_app_launcher_args(parser) +# parse the arguments +args_cli = parser.parse_args() + +# launch omniverse app +app_launcher = AppLauncher(args_cli) +simulation_app = app_launcher.app + + +"""Rest everything follows.""" + +import random +import torch + +import isaacsim.core.utils.prims as prim_utils +from isaacsim.core.api.simulation_context import SimulationContext +from isaacsim.core.cloner import GridCloner +from isaacsim.core.prims import RigidPrim +from isaacsim.core.utils.viewports import set_camera_view + +import isaaclab.sim as sim_utils +import isaaclab.terrains as terrain_gen +from isaaclab.sensors.ray_caster import MultiMeshRayCaster, MultiMeshRayCasterCfg, patterns +from isaaclab.terrains.config.rough import ROUGH_TERRAINS_CFG +from isaaclab.terrains.terrain_importer import TerrainImporter +from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR +from isaaclab.utils.math import quat_from_euler_xyz +from isaaclab.utils.timer import Timer + + +def design_scene(sim: SimulationContext, num_envs: int = 2048): + """Design the scene.""" + # Create interface to clone the scene + cloner = GridCloner(spacing=10.0) + cloner.define_base_env("/World/envs") + # Everything under the namespace "/World/envs/env_0" will be cloned + prim_utils.define_prim("/World/envs/env_0") + # Define the scene + # -- Light + cfg = sim_utils.DistantLightCfg(intensity=2000) + cfg.func("/World/light", cfg) + # -- Balls + cfg = sim_utils.SphereCfg( + radius=0.25, + rigid_props=sim_utils.RigidBodyPropertiesCfg(), + mass_props=sim_utils.MassPropertiesCfg(mass=0.5), + collision_props=sim_utils.CollisionPropertiesCfg(), + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.0, 0.0, 1.0)), + ) + cfg.func("/World/envs/env_0/ball", cfg, translation=(0.0, 0.0, 5.0)) + + for i in range(args_cli.num_objects): + object = sim_utils.CuboidCfg( + size=(0.5 + random.random() * 0.5, 0.5 + random.random() * 0.5, 0.1 + random.random() * 0.05), + rigid_props=sim_utils.RigidBodyPropertiesCfg(), + mass_props=sim_utils.MassPropertiesCfg(mass=0.5), + collision_props=sim_utils.CollisionPropertiesCfg(), + visual_material=sim_utils.PreviewSurfaceCfg( + diffuse_color=(0.0 + i / args_cli.num_objects, 0.0, 1.0 - i / args_cli.num_objects) + ), + ) + object.func( + f"/World/envs/env_0/object_{i}", + object, + translation=(0.0 + random.random(), 0.0 + random.random(), 1.0), + orientation=quat_from_euler_xyz(torch.Tensor(0), torch.Tensor(0), torch.rand(1) * torch.pi).numpy(), + ) + + # Clone the scene + cloner.define_base_env("/World/envs") + envs_prim_paths = cloner.generate_paths("/World/envs/env", num_paths=num_envs) + cloner.clone(source_prim_path="/World/envs/env_0", prim_paths=envs_prim_paths, replicate_physics=True) + physics_scene_path = sim.get_physics_context().prim_path + cloner.filter_collisions( + physics_scene_path, "/World/collisions", prim_paths=envs_prim_paths, global_paths=["/World/ground"] + ) + + +def main(): + """Main function.""" + + # Load kit helper + sim_params = { + "use_gpu": True, + "use_gpu_pipeline": True, + "use_flatcache": True, # deprecated from Isaac Sim 2023.1 onwards + "use_fabric": True, # used from Isaac Sim 2023.1 onwards + "enable_scene_query_support": True, + } + sim = SimulationContext( + physics_dt=1.0 / 60.0, rendering_dt=1.0 / 60.0, sim_params=sim_params, backend="torch", device="cuda:0" + ) + # Set main camera + set_camera_view([0.0, 30.0, 25.0], [0.0, 0.0, -2.5]) + + # Parameters + num_envs = args_cli.num_envs + # Design the scene + design_scene(sim=sim, num_envs=num_envs) + # Handler for terrains importing + terrain_importer_cfg = terrain_gen.TerrainImporterCfg( + prim_path="/World/ground", + terrain_type=args_cli.terrain_type, + terrain_generator=ROUGH_TERRAINS_CFG, + usd_path=f"{ISAAC_NUCLEUS_DIR}/Environments/Terrains/rough_plane.usd", + max_init_terrain_level=0, + num_envs=1, + ) + _ = TerrainImporter(terrain_importer_cfg) + + mesh_targets: list[MultiMeshRayCasterCfg.RaycastTargetCfg] = [ + MultiMeshRayCasterCfg.RaycastTargetCfg(target_prim_expr="/World/ground", track_mesh_transforms=False), + ] + if args_cli.num_objects != 0: + mesh_targets.append( + MultiMeshRayCasterCfg.RaycastTargetCfg( + target_prim_expr="/World/envs/env_.*/object_.*", track_mesh_transforms=True + ) + ) + # Create a ray-caster sensor + ray_caster_cfg = MultiMeshRayCasterCfg( + prim_path="/World/envs/env_.*/ball", + mesh_prim_paths=mesh_targets, + pattern_cfg=patterns.GridPatternCfg(resolution=0.1, size=(1.6, 1.0)), + attach_yaw_only=True, + debug_vis=not args_cli.headless, + ) + ray_caster = MultiMeshRayCaster(cfg=ray_caster_cfg) + # Create a view over all the balls + ball_view = RigidPrim("/World/envs/env_.*/ball", reset_xform_properties=False) + + # Play simulator + sim.reset() + + # Initialize the views + # -- balls + ball_view.initialize() + # Print the sensor information + print(ray_caster) + + # Get the initial positions of the balls + ball_initial_positions, ball_initial_orientations = ball_view.get_world_poses() + ball_initial_velocities = ball_view.get_velocities() + + # Create a counter for resetting the scene + step_count = 0 + # Simulate physics + while simulation_app.is_running(): + # If simulation is stopped, then exit. + if sim.is_stopped(): + break + # If simulation is paused, then skip. + if not sim.is_playing(): + sim.step(render=False) + continue + # Reset the scene + if step_count % 500 == 0: + # sample random indices to reset + reset_indices = torch.randint(0, num_envs, (num_envs // 2,)) + # reset the balls + ball_view.set_world_poses( + ball_initial_positions[reset_indices], ball_initial_orientations[reset_indices], indices=reset_indices + ) + ball_view.set_velocities(ball_initial_velocities[reset_indices], indices=reset_indices) + # reset the sensor + ray_caster.reset(reset_indices) + # reset the counter + step_count = 0 + # Step simulation + sim.step() + # Update the ray-caster + with Timer(f"Ray-caster update with {num_envs} x {ray_caster.num_rays} rays"): + ray_caster.update(dt=sim.get_physics_dt(), force_recompute=True) + # Update counter + step_count += 1 + + +if __name__ == "__main__": + # run the main function + main() + # close sim app + simulation_app.close() diff --git a/source/isaaclab/test/sensors/test_multi_mesh_ray_caster.py b/source/isaaclab/test/sensors/test_multi_mesh_ray_caster.py new file mode 100644 index 00000000000..60dceb43a6b --- /dev/null +++ b/source/isaaclab/test/sensors/test_multi_mesh_ray_caster.py @@ -0,0 +1,251 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + + +from __future__ import annotations + +from isaaclab.app import AppLauncher + +# launch omniverse app. Used for warp. +app_launcher = AppLauncher(headless=True) + +import numpy as np +import torch +import trimesh + +import pytest +import warp as wp + +from isaaclab.utils.math import matrix_from_quat, quat_from_euler_xyz, random_orientation +from isaaclab.utils.warp.ops import convert_to_warp_mesh, raycast_dynamic_meshes, raycast_single_mesh + + +@pytest.fixture(scope="module") +def device(): + return "cuda" if torch.cuda.is_available() else "cpu" + + +@pytest.fixture +def rays(device): + ray_starts = torch.tensor([[0, -0.35, -5], [0.25, 0.35, -5]], dtype=torch.float32, device=device).unsqueeze(0) + ray_directions = torch.tensor([[0, 0, 1], [0, 0, 1]], dtype=torch.float32, device=device).unsqueeze(0) + expected_ray_hits = torch.tensor( + [[0, -0.35, -0.5], [0.25, 0.35, -0.5]], dtype=torch.float32, device=device + ).unsqueeze(0) + return ray_starts, ray_directions, expected_ray_hits + + +@pytest.fixture +def trimesh_box(): + return trimesh.creation.box([2, 2, 1]) + + +@pytest.fixture +def single_mesh(trimesh_box, device): + wp_mesh = convert_to_warp_mesh(trimesh_box.vertices, trimesh_box.faces, device) + return wp_mesh, wp_mesh.id + + +def test_raycast_multi_cubes(device, trimesh_box, rays): + """Test raycasting against two cubes.""" + ray_starts, ray_directions, _ = rays + + trimesh_1 = trimesh_box.copy() + wp_mesh_1 = convert_to_warp_mesh(trimesh_1.vertices, trimesh_1.faces, device) + + translation = np.eye(4) + translation[:3, 3] = [0, 2, 0] + trimesh_2 = trimesh_box.copy().apply_transform(translation) + wp_mesh_2 = convert_to_warp_mesh(trimesh_2.vertices, trimesh_2.faces, device) + + # get mesh id array + mesh_ids_wp = wp.array2d([[wp_mesh_1.id, wp_mesh_2.id]], dtype=wp.uint64, device=device) + + # Static positions (no transforms passed) + ray_start = torch.tensor([[0, 0, -5], [0, 2.5, -5]], dtype=torch.float32, device=device).unsqueeze(0) + ray_hits, ray_distance, ray_normal, ray_face_id, mesh_ids = raycast_dynamic_meshes( + ray_start, + ray_directions, + mesh_ids_wp, + return_distance=True, + return_normal=True, + return_face_id=True, + return_mesh_id=True, + ) + + torch.testing.assert_close( + ray_hits, torch.tensor([[[0, 0, -0.5], [0, 2.5, -0.5]]], dtype=torch.float32, device=device) + ) + torch.testing.assert_close(ray_distance, torch.tensor([[4.5, 4.5]], dtype=torch.float32, device=device)) + torch.testing.assert_close(ray_normal, torch.tensor([[[0, 0, -1], [0, 0, -1]]], dtype=torch.float32, device=device)) + assert torch.equal(mesh_ids, torch.tensor([[0, 1]], dtype=torch.int32, device=device)) + + # Dynamic positions/orientations + ray_start = torch.tensor([[0, 0, -5], [0, 4.5, -5]], dtype=torch.float32, device=device).unsqueeze(0) + ray_hits, ray_distance, ray_normal, ray_face_id, mesh_ids = raycast_dynamic_meshes( + ray_start, + ray_directions, + mesh_ids_wp, + return_distance=True, + return_normal=True, + return_face_id=True, + mesh_positions_w=torch.tensor([[[0, 0, 0], [0, 2, 0]]], dtype=torch.float32, device=device), + mesh_orientations_w=torch.tensor([[[1, 0, 0, 0], [1, 0, 0, 0]]], dtype=torch.float32, device=device), + return_mesh_id=True, + ) + + torch.testing.assert_close( + ray_hits, torch.tensor([[[0, 0, -0.5], [0, 4.5, -0.5]]], dtype=torch.float32, device=device) + ) + torch.testing.assert_close(ray_distance, torch.tensor([[4.5, 4.5]], dtype=torch.float32, device=device)) + torch.testing.assert_close(ray_normal, torch.tensor([[[0, 0, -1], [0, 0, -1]]], dtype=torch.float32, device=device)) + assert torch.equal(mesh_ids, torch.tensor([[0, 1]], dtype=torch.int32, device=device)) + + +def test_raycast_single_cube(device, single_mesh, rays): + """Test raycasting against a single cube.""" + ray_starts, ray_directions, expected_ray_hits = rays + _, single_mesh_id = single_mesh + + ray_hits, ray_distance, ray_normal, ray_face_id = raycast_single_mesh( + ray_starts, + ray_directions, + single_mesh_id, + return_distance=True, + return_normal=True, + return_face_id=True, + ) + torch.testing.assert_close(ray_hits, expected_ray_hits) + torch.testing.assert_close(ray_distance, torch.tensor([[4.5, 4.5]], dtype=torch.float32, device=device)) + torch.testing.assert_close( + ray_normal, + torch.tensor([[[0, 0, -1], [0, 0, -1]]], dtype=torch.float32, device=device), + ) + torch.testing.assert_close(ray_face_id, torch.tensor([[3, 8]], dtype=torch.int32, device=device)) + + # check multiple meshes implementation + ray_hits, ray_distance, ray_normal, ray_face_id, _ = raycast_dynamic_meshes( + ray_starts, + ray_directions, + wp.array2d([[single_mesh_id]], dtype=wp.uint64, device=device), + return_distance=True, + return_normal=True, + return_face_id=True, + ) + torch.testing.assert_close(ray_hits, expected_ray_hits) + torch.testing.assert_close(ray_distance, torch.tensor([[4.5, 4.5]], dtype=torch.float32, device=device)) + torch.testing.assert_close(ray_normal, torch.tensor([[[0, 0, -1], [0, 0, -1]]], dtype=torch.float32, device=device)) + torch.testing.assert_close(ray_face_id, torch.tensor([[3, 8]], dtype=torch.int32, device=device)) + + +@pytest.mark.parametrize("num_samples", [10]) +def test_raycast_moving_cube(device, single_mesh, rays, num_samples): + r"""Test raycasting against a single cube with different distances. + |-------------| + |\ | + | \ | + | \ 8 | + | \ | + | \ x_1 | + | \ | + | \ | + | \ | + | \ | + | \ | + | 3 x_2 \ | + | \ | + | \| + |-------------| + + """ + ray_starts, ray_directions, expected_ray_hits = rays + _, single_mesh_id = single_mesh + + # move the cube along the z axis + for distance in torch.linspace(0, 1, num_samples, device=device): + ray_hits, ray_distance, ray_normal, ray_face_id, mesh_id = raycast_dynamic_meshes( + ray_starts, + ray_directions, + wp.array2d([[single_mesh_id]], dtype=wp.uint64, device=device), + return_distance=True, + return_normal=True, + return_face_id=True, + return_mesh_id=True, + mesh_positions_w=torch.tensor([[0, 0, distance]], dtype=torch.float32, device=device), + ) + torch.testing.assert_close( + ray_hits, + expected_ray_hits + + torch.tensor([[0, 0, distance], [0, 0, distance]], dtype=torch.float32, device=device).unsqueeze(0), + ) + torch.testing.assert_close( + ray_distance, distance + torch.tensor([[4.5, 4.5]], dtype=torch.float32, device=device) + ) + torch.testing.assert_close( + ray_normal, torch.tensor([[[0, 0, -1], [0, 0, -1]]], dtype=torch.float32, device=device) + ) + torch.testing.assert_close(ray_face_id, torch.tensor([[3, 8]], dtype=torch.int32, device=device)) + + +def test_raycast_rotated_cube(device, single_mesh, rays): + """Test raycasting against a single cube with different 90deg. orientations.""" + ray_starts, ray_directions, expected_ray_hits = rays + _, single_mesh_id = single_mesh + + cube_rotation = quat_from_euler_xyz(torch.tensor([0.0]), torch.tensor([0.0]), torch.tensor([np.pi])) + ray_hits, ray_distance, ray_normal, ray_face_id, _ = raycast_dynamic_meshes( + ray_starts, + ray_directions, + wp.array2d([[single_mesh_id]], dtype=wp.uint64, device=device), + return_distance=True, + return_normal=True, + return_face_id=True, + mesh_orientations_w=cube_rotation.unsqueeze(0), + ) + torch.testing.assert_close(ray_hits, expected_ray_hits) + torch.testing.assert_close(ray_distance, torch.tensor([[4.5, 4.5]], dtype=torch.float32, device=device)) + torch.testing.assert_close(ray_normal, torch.tensor([[[0, 0, -1], [0, 0, -1]]], dtype=torch.float32, device=device)) + # Make sure the face ids are correct. The cube is rotated by 90deg. so the face ids are different. + torch.testing.assert_close(ray_face_id, torch.tensor([[8, 3]], dtype=torch.int32, device=device)) + + +@pytest.mark.parametrize("num_random", [10]) +def test_raycast_random_cube(device, trimesh_box, single_mesh, rays, num_random): + """Test raycasting against a single cube with random poses.""" + ray_starts, ray_directions, _ = rays + _, single_mesh_id = single_mesh + + for orientation in random_orientation(num_random, device): + pos = torch.tensor([[0, 0, torch.rand(1)]], dtype=torch.float32, device=device) + tf_hom = np.eye(4) + tf_hom[:3, :3] = matrix_from_quat(orientation).cpu().numpy() + tf_hom[:3, 3] = pos.cpu().numpy() + tf_mesh = trimesh_box.copy().apply_transform(tf_hom) + + # get raycast for transformed, static mesh + wp_mesh = convert_to_warp_mesh(tf_mesh.vertices, tf_mesh.faces, device) + ray_hits, ray_distance, ray_normal, ray_face_id, _ = raycast_dynamic_meshes( + ray_starts, + ray_directions, + wp.array2d([[wp_mesh.id]], dtype=wp.uint64, device=device), + return_distance=True, + return_normal=True, + return_face_id=True, + ) + # get raycast for modified mesh + ray_hits_m, ray_distance_m, ray_normal_m, ray_face_id_m, _ = raycast_dynamic_meshes( + ray_starts, + ray_directions, + wp.array2d([[single_mesh_id]], dtype=wp.uint64, device=device), + return_distance=True, + return_normal=True, + return_face_id=True, + mesh_positions_w=pos, + mesh_orientations_w=orientation.view(1, 1, -1), + ) + torch.testing.assert_close(ray_hits, ray_hits_m) + torch.testing.assert_close(ray_distance, ray_distance_m) + torch.testing.assert_close(ray_normal, ray_normal_m) + torch.testing.assert_close(ray_face_id, ray_face_id_m) diff --git a/source/isaaclab/test/sensors/test_multi_mesh_ray_caster_camera.py b/source/isaaclab/test/sensors/test_multi_mesh_ray_caster_camera.py new file mode 100644 index 00000000000..5a3d7a5a879 --- /dev/null +++ b/source/isaaclab/test/sensors/test_multi_mesh_ray_caster_camera.py @@ -0,0 +1,863 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# ignore private usage of variables warning +# pyright: reportPrivateUsage=none + +"""Launch Isaac Sim Simulator first.""" + +from isaaclab.app import AppLauncher + +# launch omniverse app +simulation_app = AppLauncher(headless=True, enable_cameras=True).app + +"""Rest everything follows.""" + +import copy +import numpy as np +import os +import torch + +import isaacsim.core.utils.prims as prim_utils +import isaacsim.core.utils.stage as stage_utils +import omni.replicator.core as rep +import pytest +from pxr import Gf + +import isaaclab.sim as sim_utils +from isaaclab.sensors.camera import Camera, CameraCfg +from isaaclab.sensors.ray_caster import MultiMeshRayCasterCamera, MultiMeshRayCasterCameraCfg, patterns +from isaaclab.sim import PinholeCameraCfg +from isaaclab.terrains.trimesh.utils import make_plane +from isaaclab.terrains.utils import create_prim_from_mesh +from isaaclab.utils import convert_dict_to_backend +from isaaclab.utils.timer import Timer + +# sample camera poses +POSITION = [2.5, 2.5, 2.5] +QUAT_ROS = [-0.17591989, 0.33985114, 0.82047325, -0.42470819] +QUAT_OPENGL = [0.33985113, 0.17591988, 0.42470818, 0.82047324] +QUAT_WORLD = [-0.3647052, -0.27984815, -0.1159169, 0.88047623] + + +@pytest.fixture(scope="function") +def setup_simulation(): + """Fixture to set up and tear down the simulation environment.""" + # Create a new stage + stage_utils.create_new_stage() + # Simulation time-step + dt = 0.01 + # Load kit helper + sim_cfg = sim_utils.SimulationCfg(dt=dt) + sim: sim_utils.SimulationContext = sim_utils.SimulationContext(sim_cfg) + # Ground-plane + mesh = make_plane(size=(100, 100), height=0.0, center_zero=True) + create_prim_from_mesh("/World/defaultGroundPlane", mesh) + # load stage + stage_utils.update_stage() + + camera_cfg = MultiMeshRayCasterCameraCfg( + prim_path="/World/Camera", + mesh_prim_paths=["/World/defaultGroundPlane"], + update_period=0, + offset=MultiMeshRayCasterCameraCfg.OffsetCfg(pos=(0.0, 0.0, 0.0), rot=(1.0, 0.0, 0.0, 0.0), convention="world"), + debug_vis=False, + pattern_cfg=patterns.PinholeCameraPatternCfg( + focal_length=24.0, + horizontal_aperture=20.955, + height=480, + width=640, + ), + data_types=["distance_to_image_plane"], + ) + + # create xform because placement of camera directly under world is not supported + prim_utils.create_prim("/World/Camera", "Xform") + + yield sim, dt, camera_cfg + + # Cleanup + # close all the opened viewport from before. + rep.vp_manager.destroy_hydra_textures("Replicator") + # stop simulation + # note: cannot use self.sim.stop() since it does one render step after stopping!! This doesn't make sense :( + sim._timeline.stop() + # clear the stage + sim.clear_all_callbacks() + sim.clear_instance() + + +@pytest.mark.parametrize( + "convention,quat", + [ + ("ros", QUAT_ROS), + ("opengl", QUAT_OPENGL), + ("world", QUAT_WORLD), + ], +) +@pytest.mark.isaacsim_ci +def test_camera_init_offset(setup_simulation, convention, quat): + """Test camera initialization with offset using different conventions.""" + sim, dt, camera_cfg = setup_simulation + + # Create camera config with specific convention + cam_cfg_offset = copy.deepcopy(camera_cfg) + cam_cfg_offset.offset = MultiMeshRayCasterCameraCfg.OffsetCfg( + pos=POSITION, + rot=quat, + convention=convention, + ) + prim_utils.create_prim(f"/World/CameraOffset{convention.capitalize()}", "Xform") + cam_cfg_offset.prim_path = f"/World/CameraOffset{convention.capitalize()}" + + camera = MultiMeshRayCasterCamera(cam_cfg_offset) + + # play sim + sim.reset() + + # update camera + camera.update(dt) + + # check that transform is set correctly + np.testing.assert_allclose(camera.data.pos_w[0].cpu().numpy(), cam_cfg_offset.offset.pos) + + del camera + + +@pytest.mark.isaacsim_ci +def test_camera_init(setup_simulation): + """Test camera initialization.""" + sim, dt, camera_cfg = setup_simulation + + # Create camera + camera = MultiMeshRayCasterCamera(cfg=camera_cfg) + # Play sim + sim.reset() + # Check if camera is initialized + assert camera.is_initialized + # Simulate for a few steps + # note: This is a workaround to ensure that the textures are loaded. + # Check "Known Issues" section in the documentation for more details. + for _ in range(5): + sim.step() + # Check buffers that exists and have correct shapes + assert camera.data.pos_w.shape == (1, 3) + assert camera.data.quat_w_ros.shape == (1, 4) + assert camera.data.quat_w_world.shape == (1, 4) + assert camera.data.quat_w_opengl.shape == (1, 4) + assert camera.data.intrinsic_matrices.shape == (1, 3, 3) + assert camera.data.image_shape == (camera_cfg.pattern_cfg.height, camera_cfg.pattern_cfg.width) + assert camera.data.info == [{camera_cfg.data_types[0]: None}] + # Simulate physics + for _ in range(10): + # perform rendering + sim.step() + # update camera + camera.update(dt) + # check image data + for im_data in camera.data.output.values(): + assert im_data.shape == (1, camera_cfg.pattern_cfg.height, camera_cfg.pattern_cfg.width, 1) + + del camera + + +@pytest.mark.isaacsim_ci +def test_camera_resolution(setup_simulation): + """Test camera resolution is correctly set.""" + sim, dt, camera_cfg = setup_simulation + + # Create camera + camera = MultiMeshRayCasterCamera(cfg=camera_cfg) + # Play sim + sim.reset() + # Simulate for a few steps + # note: This is a workaround to ensure that the textures are loaded. + # Check "Known Issues" section in the documentation for more details. + for _ in range(5): + sim.step() + camera.update(dt) + # access image data and compare shapes + for im_data in camera.data.output.values(): + assert im_data.shape == (1, camera_cfg.pattern_cfg.height, camera_cfg.pattern_cfg.width, 1) + + del camera + + +@pytest.mark.isaacsim_ci +def test_camera_init_intrinsic_matrix(setup_simulation): + """Test camera initialization from intrinsic matrix.""" + sim, dt, camera_cfg = setup_simulation + + # get the first camera + camera_1 = MultiMeshRayCasterCamera(cfg=camera_cfg) + # get intrinsic matrix + sim.reset() + intrinsic_matrix = camera_1.data.intrinsic_matrices[0].cpu().flatten().tolist() + + # initialize from intrinsic matrix + intrinsic_camera_cfg = MultiMeshRayCasterCameraCfg( + prim_path="/World/Camera", + mesh_prim_paths=["/World/defaultGroundPlane"], + update_period=0, + offset=MultiMeshRayCasterCameraCfg.OffsetCfg(pos=(0.0, 0.0, 0.0), rot=(1.0, 0.0, 0.0, 0.0), convention="world"), + debug_vis=False, + pattern_cfg=patterns.PinholeCameraPatternCfg.from_intrinsic_matrix( + intrinsic_matrix=intrinsic_matrix, + height=camera_cfg.pattern_cfg.height, + width=camera_cfg.pattern_cfg.width, + focal_length=camera_cfg.pattern_cfg.focal_length, + ), + data_types=["distance_to_image_plane"], + ) + camera_2 = MultiMeshRayCasterCamera(cfg=intrinsic_camera_cfg) + + # play sim + sim.reset() + sim.play() + + # update cameras + camera_1.update(dt) + camera_2.update(dt) + + # check image data + torch.testing.assert_close( + camera_1.data.output["distance_to_image_plane"], + camera_2.data.output["distance_to_image_plane"], + ) + # check that both intrinsic matrices are the same + torch.testing.assert_close( + camera_1.data.intrinsic_matrices[0], + camera_2.data.intrinsic_matrices[0], + ) + + del camera_1, camera_2 + + +@pytest.mark.isaacsim_ci +def test_multi_camera_init(setup_simulation): + """Test multi-camera initialization.""" + sim, dt, camera_cfg = setup_simulation + + # -- camera 1 + cam_cfg_1 = copy.deepcopy(camera_cfg) + cam_cfg_1.prim_path = "/World/Camera_0" + prim_utils.create_prim("/World/Camera_0", "Xform") + # Create camera + cam_1 = MultiMeshRayCasterCamera(cam_cfg_1) + + # -- camera 2 + cam_cfg_2 = copy.deepcopy(camera_cfg) + cam_cfg_2.prim_path = "/World/Camera_1" + prim_utils.create_prim("/World/Camera_1", "Xform") + # Create camera + cam_2 = MultiMeshRayCasterCamera(cam_cfg_2) + + # play sim + sim.reset() + + # Simulate for a few steps + # note: This is a workaround to ensure that the textures are loaded. + # Check "Known Issues" section in the documentation for more details. + for _ in range(5): + sim.step() + # Simulate physics + for _ in range(10): + # perform rendering + sim.step() + # update camera + cam_1.update(dt) + cam_2.update(dt) + # check image data + for cam in [cam_1, cam_2]: + for im_data in cam.data.output.values(): + assert im_data.shape == (1, camera_cfg.pattern_cfg.height, camera_cfg.pattern_cfg.width, 1) + + del cam_1, cam_2 + + +@pytest.mark.isaacsim_ci +def test_camera_set_world_poses(setup_simulation): + """Test camera function to set specific world pose.""" + sim, dt, camera_cfg = setup_simulation + + camera = MultiMeshRayCasterCamera(camera_cfg) + # play sim + sim.reset() + + # convert to torch tensors + position = torch.tensor([POSITION], dtype=torch.float32, device=camera.device) + orientation = torch.tensor([QUAT_WORLD], dtype=torch.float32, device=camera.device) + # set new pose + camera.set_world_poses(position.clone(), orientation.clone(), convention="world") + + # check if transform correctly set in output + torch.testing.assert_close(camera.data.pos_w, position) + torch.testing.assert_close(camera.data.quat_w_world, orientation) + + del camera + + +@pytest.mark.isaacsim_ci +def test_camera_set_world_poses_from_view(setup_simulation): + """Test camera function to set specific world pose from view.""" + sim, dt, camera_cfg = setup_simulation + + camera = MultiMeshRayCasterCamera(camera_cfg) + # play sim + sim.reset() + + # convert to torch tensors + eyes = torch.tensor([POSITION], dtype=torch.float32, device=camera.device) + targets = torch.tensor([[0.0, 0.0, 0.0]], dtype=torch.float32, device=camera.device) + quat_ros_gt = torch.tensor([QUAT_ROS], dtype=torch.float32, device=camera.device) + # set new pose + camera.set_world_poses_from_view(eyes.clone(), targets.clone()) + + # check if transform correctly set in output + torch.testing.assert_close(camera.data.pos_w, eyes) + torch.testing.assert_close(camera.data.quat_w_ros, quat_ros_gt) + + del camera + + +@pytest.mark.parametrize("height,width", [(240, 320), (480, 640)]) +@pytest.mark.isaacsim_ci +def test_intrinsic_matrix(setup_simulation, height, width): + """Checks that the camera's set and retrieve methods work for intrinsic matrix.""" + sim, dt, camera_cfg = setup_simulation + + camera_cfg_copy = copy.deepcopy(camera_cfg) + camera_cfg_copy.pattern_cfg.height = height + camera_cfg_copy.pattern_cfg.width = width + camera = MultiMeshRayCasterCamera(camera_cfg_copy) + # play sim + sim.reset() + # Desired properties (obtained from realsense camera at 320x240 resolution) + rs_intrinsic_matrix = [229.31640625, 0.0, 164.810546875, 0.0, 229.826171875, 122.1650390625, 0.0, 0.0, 1.0] + rs_intrinsic_matrix = torch.tensor(rs_intrinsic_matrix, device=camera.device).reshape(3, 3).unsqueeze(0) + # Set matrix into simulator + camera.set_intrinsic_matrices(rs_intrinsic_matrix.clone()) + # Simulate for a few steps + # note: This is a workaround to ensure that the textures are loaded. + # Check "Known Issues" section in the documentation for more details. + for _ in range(5): + sim.step() + # Simulate physics + for _ in range(10): + # perform rendering + sim.step() + # update camera + camera.update(dt) + # Check that matrix is correct + torch.testing.assert_close(rs_intrinsic_matrix, camera.data.intrinsic_matrices) + + del camera + + +@pytest.mark.isaacsim_ci +def test_throughput(setup_simulation): + """Checks that the single camera gets created properly with a rig.""" + sim, dt, camera_cfg = setup_simulation + + # Create directory temp dir to dump the results + file_dir = os.path.dirname(os.path.realpath(__file__)) + temp_dir = os.path.join(file_dir, "output", "camera", "throughput") + os.makedirs(temp_dir, exist_ok=True) + # Create replicator writer + rep_writer = rep.BasicWriter(output_dir=temp_dir, frame_padding=3) + # create camera + camera_cfg_copy = copy.deepcopy(camera_cfg) + camera_cfg_copy.pattern_cfg.height = 480 + camera_cfg_copy.pattern_cfg.width = 640 + camera = MultiMeshRayCasterCamera(camera_cfg_copy) + + # Play simulator + sim.reset() + + # Set camera pose + eyes = torch.tensor([[2.5, 2.5, 2.5]], dtype=torch.float32, device=camera.device) + targets = torch.tensor([[0.0, 0.0, 0.0]], dtype=torch.float32, device=camera.device) + camera.set_world_poses_from_view(eyes, targets) + + # Simulate for a few steps + # note: This is a workaround to ensure that the textures are loaded. + # Check "Known Issues" section in the documentation for more details. + for _ in range(5): + sim.step() + # Simulate physics + for _ in range(5): + # perform rendering + sim.step() + # update camera + with Timer(f"Time taken for updating camera with shape {camera.image_shape}"): + camera.update(dt) + # Save images + with Timer(f"Time taken for writing data with shape {camera.image_shape} "): + # Pack data back into replicator format to save them using its writer + rep_output = {"annotators": {}} + camera_data = convert_dict_to_backend(camera.data.output, backend="numpy") + for key, data, info in zip(camera_data.keys(), camera_data.values(), camera.data.info[0].values()): + if info is not None: + rep_output["annotators"][key] = {"render_product": {"data": data, **info}} + else: + rep_output["annotators"][key] = {"render_product": {"data": data}} + # Save images + rep_output["trigger_outputs"] = {"on_time": camera.frame[0]} + rep_writer.write(rep_output) + print("----------------------------------------") + # Check image data + for im_data in camera.data.output.values(): + assert im_data.shape == (1, camera_cfg_copy.pattern_cfg.height, camera_cfg_copy.pattern_cfg.width, 1) + + del camera + + +@pytest.mark.parametrize( + "data_types", + [ + ["distance_to_image_plane", "distance_to_camera", "normals"], + ["distance_to_image_plane"], + ["distance_to_camera"], + ], +) +@pytest.mark.isaacsim_ci +def test_output_equal_to_usdcamera(setup_simulation, data_types): + """Test that ray caster camera output equals USD camera output.""" + sim, dt, camera_cfg = setup_simulation + + camera_pattern_cfg = patterns.PinholeCameraPatternCfg( + focal_length=24.0, + horizontal_aperture=20.955, + height=240, + width=320, + ) + prim_utils.create_prim("/World/Camera_warp", "Xform") + camera_cfg_warp = MultiMeshRayCasterCameraCfg( + prim_path="/World/Camera", + mesh_prim_paths=["/World/defaultGroundPlane"], + update_period=0, + offset=MultiMeshRayCasterCameraCfg.OffsetCfg(pos=(0.0, 0.0, 0.0), rot=(1.0, 0.0, 0.0, 0.0)), + debug_vis=False, + pattern_cfg=camera_pattern_cfg, + data_types=data_types, + ) + + camera_warp = MultiMeshRayCasterCamera(camera_cfg_warp) + + # create usd camera + camera_cfg_usd = CameraCfg( + height=240, + width=320, + prim_path="/World/Camera_usd", + update_period=0, + data_types=data_types, + spawn=PinholeCameraCfg( + focal_length=24.0, focus_distance=400.0, horizontal_aperture=20.955, clipping_range=(1e-4, 1.0e5) + ), + ) + camera_usd = Camera(camera_cfg_usd) + + # play sim + sim.reset() + sim.play() + + # convert to torch tensors + eyes = torch.tensor([[2.5, 2.5, 4.5]], dtype=torch.float32, device=camera_warp.device) + targets = torch.tensor([[0.0, 0.0, 0.0]], dtype=torch.float32, device=camera_warp.device) + # set views + camera_warp.set_world_poses_from_view(eyes, targets) + camera_usd.set_world_poses_from_view(eyes, targets) + + # perform steps + for _ in range(5): + sim.step() + + # update camera + camera_usd.update(dt) + camera_warp.update(dt) + + # check the intrinsic matrices + torch.testing.assert_close( + camera_usd.data.intrinsic_matrices, + camera_warp.data.intrinsic_matrices, + ) + + # check the apertures + torch.testing.assert_close( + camera_usd._sensor_prims[0].GetHorizontalApertureAttr().Get(), + camera_cfg_warp.pattern_cfg.horizontal_aperture, + ) + + # check image data + for data_type in data_types: + if data_type in camera_usd.data.output and data_type in camera_warp.data.output: + if data_type == "distance_to_camera": + torch.testing.assert_close( + camera_usd.data.output[data_type], + camera_warp.data.output[data_type], + atol=5e-5, + rtol=5e-6, + ) + elif data_type == "normals": + # NOTE: floating point issues of ~1e-5, so using atol and rtol in this case + torch.testing.assert_close( + camera_usd.data.output[data_type][..., :3], + camera_warp.data.output[data_type], + rtol=1e-5, + atol=1e-4, + ) + else: + torch.testing.assert_close( + camera_usd.data.output[data_type], + camera_warp.data.output[data_type], + ) + + del camera_usd, camera_warp + + +@pytest.mark.isaacsim_ci +def test_output_equal_to_usdcamera_offset(setup_simulation): + """Test that ray caster camera output equals USD camera output with offset.""" + sim, dt, camera_cfg = setup_simulation + offset_rot = (-0.1251, 0.3617, 0.8731, -0.3020) + + camera_pattern_cfg = patterns.PinholeCameraPatternCfg( + focal_length=24.0, + horizontal_aperture=20.955, + height=240, + width=320, + ) + prim_utils.create_prim("/World/Camera_warp", "Xform") + camera_cfg_warp = MultiMeshRayCasterCameraCfg( + prim_path="/World/Camera", + mesh_prim_paths=["/World/defaultGroundPlane"], + update_period=0, + offset=MultiMeshRayCasterCameraCfg.OffsetCfg(pos=(2.5, 2.5, 4.0), rot=offset_rot, convention="ros"), + debug_vis=False, + pattern_cfg=camera_pattern_cfg, + data_types=["distance_to_image_plane", "distance_to_camera", "normals"], + ) + camera_warp = MultiMeshRayCasterCamera(camera_cfg_warp) + + # create usd camera + camera_cfg_usd = CameraCfg( + height=240, + width=320, + prim_path="/World/Camera_usd", + update_period=0, + data_types=["distance_to_image_plane", "distance_to_camera", "normals"], + spawn=PinholeCameraCfg( + focal_length=24.0, focus_distance=400.0, horizontal_aperture=20.955, clipping_range=(1e-6, 1.0e5) + ), + offset=CameraCfg.OffsetCfg(pos=(2.5, 2.5, 4.0), rot=offset_rot, convention="ros"), + ) + camera_usd = Camera(camera_cfg_usd) + + # play sim + sim.reset() + sim.play() + + # perform steps + for _ in range(5): + sim.step() + + # update camera + camera_usd.update(dt) + camera_warp.update(dt) + + # check image data + torch.testing.assert_close( + camera_usd.data.output["distance_to_image_plane"], + camera_warp.data.output["distance_to_image_plane"], + atol=5e-5, + rtol=5e-6, + ) + torch.testing.assert_close( + camera_usd.data.output["distance_to_camera"], + camera_warp.data.output["distance_to_camera"], + atol=5e-5, + rtol=5e-6, + ) + + # check normals + # NOTE: floating point issues of ~1e-5, so using atol and rtol in this case + torch.testing.assert_close( + camera_usd.data.output["normals"][..., :3], + camera_warp.data.output["normals"], + rtol=1e-5, + atol=1e-4, + ) + + del camera_usd, camera_warp + + +@pytest.mark.isaacsim_ci +def test_output_equal_to_usdcamera_prim_offset(setup_simulation): + """Test that the output of the ray caster camera is equal to the output of the usd camera when both are placed + under an XForm prim that is translated and rotated from the world origin.""" + sim, dt, camera_cfg = setup_simulation + + offset_rot = [-0.1251, 0.3617, 0.8731, -0.3020] + + # gf quat + gf_quatf = Gf.Quatd() + gf_quatf.SetReal(QUAT_OPENGL[0]) + gf_quatf.SetImaginary(tuple(QUAT_OPENGL[1:])) + + camera_pattern_cfg = patterns.PinholeCameraPatternCfg( + focal_length=24.0, + horizontal_aperture=20.955, + height=240, + width=320, + ) + prim_raycast_cam = prim_utils.create_prim("/World/Camera_warp", "Xform") + prim_raycast_cam.GetAttribute("xformOp:translate").Set(tuple(POSITION)) + prim_raycast_cam.GetAttribute("xformOp:orient").Set(gf_quatf) + + camera_cfg_warp = MultiMeshRayCasterCameraCfg( + prim_path="/World/Camera_warp", + mesh_prim_paths=["/World/defaultGroundPlane"], + update_period=0, + offset=MultiMeshRayCasterCameraCfg.OffsetCfg(pos=(0, 0, 2.0), rot=offset_rot, convention="ros"), + debug_vis=False, + pattern_cfg=camera_pattern_cfg, + data_types=["distance_to_image_plane", "distance_to_camera", "normals"], + ) + + camera_warp = MultiMeshRayCasterCamera(camera_cfg_warp) + + # create usd camera + camera_cfg_usd = CameraCfg( + height=240, + width=320, + prim_path="/World/Camera_usd/camera", + update_period=0, + data_types=["distance_to_image_plane", "distance_to_camera", "normals"], + spawn=PinholeCameraCfg( + focal_length=24.0, focus_distance=400.0, horizontal_aperture=20.955, clipping_range=(1e-6, 1.0e5) + ), + offset=CameraCfg.OffsetCfg(pos=(0, 0, 2.0), rot=offset_rot, convention="ros"), + update_latest_camera_pose=True, + ) + prim_usd = prim_utils.create_prim("/World/Camera_usd", "Xform") + prim_usd.GetAttribute("xformOp:translate").Set(tuple(POSITION)) + prim_usd.GetAttribute("xformOp:orient").Set(gf_quatf) + + camera_usd = Camera(camera_cfg_usd) + + # play sim + sim.reset() + sim.play() + + # perform steps + for _ in range(5): + sim.step() + + # update camera + camera_usd.update(dt) + camera_warp.update(dt) + + # check if pos and orientation are correct + torch.testing.assert_close(camera_warp.data.pos_w[0], camera_usd.data.pos_w[0]) + torch.testing.assert_close(camera_warp.data.quat_w_ros[0], camera_usd.data.quat_w_ros[0]) + + # check image data + torch.testing.assert_close( + camera_usd.data.output["distance_to_image_plane"], + camera_warp.data.output["distance_to_image_plane"], + atol=5e-5, + rtol=5e-6, + ) + torch.testing.assert_close( + camera_usd.data.output["distance_to_camera"], + camera_warp.data.output["distance_to_camera"], + rtol=4e-6, + atol=2e-5, + ) + + # check normals + # NOTE: floating point issues of ~1e-5, so using atol and rtol in this case + torch.testing.assert_close( + camera_usd.data.output["normals"][..., :3], + camera_warp.data.output["normals"], + rtol=1e-5, + atol=1e-4, + ) + + del camera_usd, camera_warp + + +@pytest.mark.parametrize("height,width", [(540, 960), (240, 320)]) +@pytest.mark.isaacsim_ci +def test_output_equal_to_usd_camera_intrinsics(setup_simulation, height, width): + """Test that the output of the ray caster camera and usd camera are the same when both are + initialized with the same intrinsic matrix.""" + sim, dt, camera_cfg = setup_simulation + + # create cameras + offset_rot = [-0.1251, 0.3617, 0.8731, -0.3020] + offset_pos = (2.5, 2.5, 4.0) + intrinsics = [380.0831, 0.0, width / 2, 0.0, 380.0831, height / 2, 0.0, 0.0, 1.0] + prim_utils.create_prim("/World/Camera_warp", "Xform") + # get camera cfgs + camera_warp_cfg = MultiMeshRayCasterCameraCfg( + prim_path="/World/Camera_warp", + mesh_prim_paths=["/World/defaultGroundPlane"], + offset=MultiMeshRayCasterCameraCfg.OffsetCfg(pos=offset_pos, rot=offset_rot, convention="ros"), + debug_vis=False, + pattern_cfg=patterns.PinholeCameraPatternCfg.from_intrinsic_matrix( + intrinsic_matrix=intrinsics, + height=height, + width=width, + focal_length=38.0, + ), + max_distance=25.0, + data_types=["distance_to_image_plane"], + ) + camera_usd_cfg = CameraCfg( + prim_path="/World/Camera_usd", + offset=CameraCfg.OffsetCfg(pos=offset_pos, rot=offset_rot, convention="ros"), + spawn=PinholeCameraCfg.from_intrinsic_matrix( + intrinsic_matrix=intrinsics, + height=height, + width=width, + clipping_range=(0.01, 25), + focal_length=38.0, + ), + height=height, + width=width, + data_types=["distance_to_image_plane"], + ) + + # set aperture offsets to 0, as currently not supported for usd camera + camera_warp_cfg.pattern_cfg.horizontal_aperture_offset = 0 + camera_warp_cfg.pattern_cfg.vertical_aperture_offset = 0 + camera_usd_cfg.spawn.horizontal_aperture_offset = 0 + camera_usd_cfg.spawn.vertical_aperture_offset = 0 + # init cameras + camera_warp = MultiMeshRayCasterCamera(camera_warp_cfg) + camera_usd = Camera(camera_usd_cfg) + + # play sim + sim.reset() + sim.play() + + # perform steps + for _ in range(5): + sim.step() + + # update camera + camera_usd.update(dt) + camera_warp.update(dt) + + # filter nan and inf from output + cam_warp_output = camera_warp.data.output["distance_to_image_plane"].clone() + cam_usd_output = camera_usd.data.output["distance_to_image_plane"].clone() + cam_warp_output[torch.isnan(cam_warp_output)] = 0 + cam_warp_output[torch.isinf(cam_warp_output)] = 0 + cam_usd_output[torch.isnan(cam_usd_output)] = 0 + cam_usd_output[torch.isinf(cam_usd_output)] = 0 + + # check that both have the same intrinsic matrices + torch.testing.assert_close(camera_warp.data.intrinsic_matrices[0], camera_usd.data.intrinsic_matrices[0]) + + # check the apertures + torch.testing.assert_close( + camera_usd._sensor_prims[0].GetHorizontalApertureAttr().Get(), + camera_warp_cfg.pattern_cfg.horizontal_aperture, + ) + torch.testing.assert_close( + camera_usd._sensor_prims[0].GetVerticalApertureAttr().Get(), + camera_warp_cfg.pattern_cfg.vertical_aperture, + ) + + # check image data + torch.testing.assert_close( + cam_warp_output, + cam_usd_output, + atol=5e-5, + rtol=5e-6, + ) + + del camera_usd, camera_warp + + +@pytest.mark.isaacsim_ci +def test_output_equal_to_usd_camera_when_intrinsics_set(setup_simulation): + """Test that the output of the ray caster camera is equal to the output of the usd camera when both are placed + under an XForm prim and an intrinsic matrix is set.""" + sim, dt, camera_cfg = setup_simulation + + camera_pattern_cfg = patterns.PinholeCameraPatternCfg( + focal_length=24.0, + horizontal_aperture=20.955, + height=540, + width=960, + ) + camera_cfg_warp = MultiMeshRayCasterCameraCfg( + prim_path="/World/Camera", + mesh_prim_paths=["/World/defaultGroundPlane"], + update_period=0, + offset=MultiMeshRayCasterCameraCfg.OffsetCfg(pos=(0.0, 0.0, 0.0), rot=(1.0, 0.0, 0.0, 0.0)), + debug_vis=False, + pattern_cfg=camera_pattern_cfg, + data_types=["distance_to_camera"], + ) + + camera_warp = MultiMeshRayCasterCamera(camera_cfg_warp) + + # create usd camera + camera_cfg_usd = CameraCfg( + height=540, + width=960, + prim_path="/World/Camera_usd", + update_period=0, + data_types=["distance_to_camera"], + spawn=PinholeCameraCfg( + focal_length=24.0, focus_distance=400.0, horizontal_aperture=20.955, clipping_range=(1e-4, 1.0e5) + ), + ) + camera_usd = Camera(camera_cfg_usd) + + # play sim + sim.reset() + sim.play() + + # set intrinsic matrix + # NOTE: extend the test to cover aperture offsets once supported by the usd camera + intrinsic_matrix = torch.tensor( + [[380.0831, 0.0, camera_cfg_usd.width / 2, 0.0, 380.0831, camera_cfg_usd.height / 2, 0.0, 0.0, 1.0]], + device=camera_warp.device, + ).reshape(1, 3, 3) + camera_warp.set_intrinsic_matrices(intrinsic_matrix, focal_length=10) + camera_usd.set_intrinsic_matrices(intrinsic_matrix, focal_length=10) + + # set camera position + camera_warp.set_world_poses_from_view( + eyes=torch.tensor([[0.0, 0.0, 5.0]], device=camera_warp.device), + targets=torch.tensor([[0.0, 0.0, 0.0]], device=camera_warp.device), + ) + camera_usd.set_world_poses_from_view( + eyes=torch.tensor([[0.0, 0.0, 5.0]], device=camera_usd.device), + targets=torch.tensor([[0.0, 0.0, 0.0]], device=camera_usd.device), + ) + + # perform steps + for _ in range(5): + sim.step() + + # update camera + camera_usd.update(dt) + camera_warp.update(dt) + + # check image data + torch.testing.assert_close( + camera_usd.data.output["distance_to_camera"], + camera_warp.data.output["distance_to_camera"], + rtol=5e-3, + atol=1e-4, + ) + + del camera_usd, camera_warp diff --git a/source/isaaclab/test/sensors/test_ray_caster.py b/source/isaaclab/test/sensors/test_ray_caster.py new file mode 100644 index 00000000000..1a2f34af1c8 --- /dev/null +++ b/source/isaaclab/test/sensors/test_ray_caster.py @@ -0,0 +1,240 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import numpy as np +import torch +import trimesh + +import pytest + +from isaaclab.app import AppLauncher + +# launch omniverse app +simulation_app = AppLauncher(headless=True, enable_cameras=True).app + +# Import after app launch +import warp as wp + +from isaaclab.utils.math import matrix_from_quat, quat_from_euler_xyz, random_orientation +from isaaclab.utils.warp.ops import convert_to_warp_mesh, raycast_dynamic_meshes, raycast_mesh + + +@pytest.fixture(scope="module") +def raycast_setup(): + device = "cuda" if torch.cuda.is_available() else "cpu" + # Base trimesh cube and its Warp conversion + trimesh_mesh = trimesh.creation.box([2, 2, 1]) + single_mesh = [ + convert_to_warp_mesh( + trimesh_mesh.vertices, + trimesh_mesh.faces, + device, + ) + ] + single_mesh_id = single_mesh[0].id + + # Rays + ray_starts = torch.tensor([[0, -0.35, -5], [0.25, 0.35, -5]], dtype=torch.float32, device=device).unsqueeze(0) + ray_directions = torch.tensor([[0, 0, 1], [0, 0, 1]], dtype=torch.float32, device=device).unsqueeze(0) + expected_ray_hits = torch.tensor( + [[0, -0.35, -0.5], [0.25, 0.35, -0.5]], dtype=torch.float32, device=device + ).unsqueeze(0) + + return { + "device": device, + "trimesh_mesh": trimesh_mesh, + "single_mesh_id": single_mesh_id, + "wp_mesh": single_mesh[0], + "ray_starts": ray_starts, + "ray_directions": ray_directions, + "expected_ray_hits": expected_ray_hits, + } + + +def test_raycast_multi_cubes(raycast_setup): + device = raycast_setup["device"] + base_tm = raycast_setup["trimesh_mesh"] + + tm1 = base_tm.copy() + wp_mesh_1 = convert_to_warp_mesh(tm1.vertices, tm1.faces, device) + + translation = np.eye(4) + translation[:3, 3] = [0, 2, 0] + tm2 = base_tm.copy().apply_transform(translation) + wp_mesh_2 = convert_to_warp_mesh(tm2.vertices, tm2.faces, device) + + mesh_ids_wp = wp.array2d([[wp_mesh_1.id, wp_mesh_2.id]], dtype=wp.uint64, device=device) + + ray_directions = raycast_setup["ray_directions"] + + # Case 1 + ray_start = torch.tensor([[0, 0, -5], [0, 2.5, -5]], dtype=torch.float32, device=device).unsqueeze(0) + ray_hits, ray_distance, ray_normal, ray_face_id, mesh_ids = raycast_dynamic_meshes( + ray_start, + ray_directions, + mesh_ids_wp, + return_distance=True, + return_normal=True, + return_face_id=True, + return_mesh_id=True, + ) + + torch.testing.assert_close(ray_hits, torch.tensor([[[0, 0, -0.5], [0, 2.5, -0.5]]], device=device)) + torch.testing.assert_close(ray_distance, torch.tensor([[4.5, 4.5]], device=device)) + torch.testing.assert_close(ray_normal, torch.tensor([[[0, 0, -1], [0, 0, -1]]], device=device, dtype=torch.float32)) + assert torch.equal(mesh_ids, torch.tensor([[0, 1]], dtype=torch.int32, device=device)) + + # Case 2 (explicit poses/orientations) + ray_start = torch.tensor([[0, 0, -5], [0, 4.5, -5]], dtype=torch.float32, device=device).unsqueeze(0) + ray_hits, ray_distance, ray_normal, ray_face_id, mesh_ids = raycast_dynamic_meshes( + ray_start, + ray_directions, + mesh_ids_wp, + return_distance=True, + return_normal=True, + return_face_id=True, + mesh_positions_w=torch.tensor([[[0, 0, 0], [0, 2, 0]]], dtype=torch.float32, device=device), + mesh_orientations_w=torch.tensor([[[1, 0, 0, 0], [1, 0, 0, 0]]], dtype=torch.float32, device=device), + return_mesh_id=True, + ) + + torch.testing.assert_close(ray_hits, torch.tensor([[[0, 0, -0.5], [0, 4.5, -0.5]]], device=device)) + torch.testing.assert_close(ray_distance, torch.tensor([[4.5, 4.5]], device=device)) + torch.testing.assert_close(ray_normal, torch.tensor([[[0, 0, -1], [0, 0, -1]]], device=device, dtype=torch.float32)) + assert torch.equal(mesh_ids, torch.tensor([[0, 1]], dtype=torch.int32, device=device)) + + +def test_raycast_single_cube(raycast_setup): + device = raycast_setup["device"] + ray_starts = raycast_setup["ray_starts"] + ray_directions = raycast_setup["ray_directions"] + mesh = raycast_setup["wp_mesh"] + expected_ray_hits = raycast_setup["expected_ray_hits"] + single_mesh_id = raycast_setup["single_mesh_id"] + + # Single-mesh helper + ray_hits, ray_distance, ray_normal, ray_face_id = raycast_mesh( + ray_starts, + ray_directions, + mesh, + return_distance=True, + return_normal=True, + return_face_id=True, + ) + torch.testing.assert_close(ray_hits, expected_ray_hits) + torch.testing.assert_close(ray_distance, torch.tensor([[4.5, 4.5]], device=device)) + torch.testing.assert_close(ray_normal, torch.tensor([[[0, 0, -1], [0, 0, -1]]], device=device, dtype=torch.float32)) + torch.testing.assert_close(ray_face_id, torch.tensor([[3, 8]], dtype=torch.int32, device=device)) + + # Multi-mesh API with one mesh + ray_hits, ray_distance, ray_normal, ray_face_id, _ = raycast_dynamic_meshes( + ray_starts, + ray_directions, + wp.array2d([[single_mesh_id]], dtype=wp.uint64, device=device), + return_distance=True, + return_normal=True, + return_face_id=True, + ) + torch.testing.assert_close(ray_hits, expected_ray_hits) + torch.testing.assert_close(ray_distance, torch.tensor([[4.5, 4.5]], device=device)) + torch.testing.assert_close(ray_normal, torch.tensor([[[0, 0, -1], [0, 0, -1]]], device=device, dtype=torch.float32)) + torch.testing.assert_close(ray_face_id, torch.tensor([[3, 8]], dtype=torch.int32, device=device)) + + +def test_raycast_moving_cube(raycast_setup): + device = raycast_setup["device"] + ray_starts = raycast_setup["ray_starts"] + ray_directions = raycast_setup["ray_directions"] + single_mesh_id = raycast_setup["single_mesh_id"] + expected_ray_hits = raycast_setup["expected_ray_hits"] + + for distance in torch.linspace(0, 1, 10, device=device): + ray_hits, ray_distance, ray_normal, ray_face_id, mesh_id = raycast_dynamic_meshes( + ray_starts, + ray_directions, + wp.array2d([[single_mesh_id]], dtype=wp.uint64, device=device), + return_distance=True, + return_normal=True, + return_face_id=True, + return_mesh_id=True, + mesh_positions_w=torch.tensor([[0, 0, distance.item()]], dtype=torch.float32, device=device), + ) + offset = torch.tensor([[0, 0, distance.item()], [0, 0, distance.item()]], dtype=torch.float32, device=device) + torch.testing.assert_close(ray_hits, expected_ray_hits + offset.unsqueeze(0)) + torch.testing.assert_close(ray_distance, distance + torch.tensor([[4.5, 4.5]], device=device)) + torch.testing.assert_close( + ray_normal, torch.tensor([[[0, 0, -1], [0, 0, -1]]], device=device, dtype=torch.float32) + ) + torch.testing.assert_close(ray_face_id, torch.tensor([[3, 8]], dtype=torch.int32, device=device)) + + +def test_raycast_rotated_cube(raycast_setup): + device = raycast_setup["device"] + ray_starts = raycast_setup["ray_starts"] + ray_directions = raycast_setup["ray_directions"] + single_mesh_id = raycast_setup["single_mesh_id"] + expected_ray_hits = raycast_setup["expected_ray_hits"] + + cube_rotation = quat_from_euler_xyz(torch.tensor([0.0]), torch.tensor([0.0]), torch.tensor([np.pi])) + ray_hits, ray_distance, ray_normal, ray_face_id, _ = raycast_dynamic_meshes( + ray_starts, + ray_directions, + wp.array2d([[single_mesh_id]], dtype=wp.uint64, device=device), + return_distance=True, + return_normal=True, + return_face_id=True, + mesh_orientations_w=cube_rotation.unsqueeze(0), + ) + torch.testing.assert_close(ray_hits, expected_ray_hits) + torch.testing.assert_close(ray_distance, torch.tensor([[4.5, 4.5]], device=device)) + torch.testing.assert_close(ray_normal, torch.tensor([[[0, 0, -1], [0, 0, -1]]], device=device, dtype=torch.float32)) + # Rotated cube swaps face IDs + torch.testing.assert_close(ray_face_id, torch.tensor([[8, 3]], dtype=torch.int32, device=device)) + + +def test_raycast_random_cube(raycast_setup): + device = raycast_setup["device"] + base_tm = raycast_setup["trimesh_mesh"] + ray_starts = raycast_setup["ray_starts"] + ray_directions = raycast_setup["ray_directions"] + single_mesh_id = raycast_setup["single_mesh_id"] + + for orientation in random_orientation(10, device): + pos = torch.tensor([[0.0, 0.0, torch.rand(1, device=device).item()]], dtype=torch.float32, device=device) + + tf_hom = np.eye(4) + tf_hom[:3, :3] = matrix_from_quat(orientation).cpu().numpy() + tf_hom[:3, 3] = pos.squeeze(0).cpu().numpy() + + tf_mesh = base_tm.copy().apply_transform(tf_hom) + wp_mesh = convert_to_warp_mesh(tf_mesh.vertices, tf_mesh.faces, device) + + # Raycast transformed, static mesh + ray_hits, ray_distance, ray_normal, ray_face_id, _ = raycast_dynamic_meshes( + ray_starts, + ray_directions, + wp.array2d([[wp_mesh.id]], dtype=wp.uint64, device=device), + return_distance=True, + return_normal=True, + return_face_id=True, + ) + # Raycast original mesh with pose provided + ray_hits_m, ray_distance_m, ray_normal_m, ray_face_id_m, _ = raycast_dynamic_meshes( + ray_starts, + ray_directions, + wp.array2d([[single_mesh_id]], dtype=wp.uint64, device=device), + return_distance=True, + return_normal=True, + return_face_id=True, + mesh_positions_w=pos, + mesh_orientations_w=orientation.view(1, 1, -1), + ) + + torch.testing.assert_close(ray_hits, ray_hits_m) + torch.testing.assert_close(ray_distance, ray_distance_m) + torch.testing.assert_close(ray_normal, ray_normal_m) + torch.testing.assert_close(ray_face_id, ray_face_id_m) diff --git a/source/isaaclab/test/sim/test_utils.py b/source/isaaclab/test/sim/test_utils.py index 62188d8a73a..5332e155939 100644 --- a/source/isaaclab/test/sim/test_utils.py +++ b/source/isaaclab/test/sim/test_utils.py @@ -13,6 +13,7 @@ """Rest everything follows.""" import numpy as np +import torch import isaacsim.core.utils.prims as prim_utils import isaacsim.core.utils.stage as stage_utils @@ -20,6 +21,7 @@ from pxr import Sdf, Usd, UsdGeom, UsdPhysics import isaaclab.sim as sim_utils +import isaaclab.utils.math as math_utils from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR, ISAACLAB_NUCLEUS_DIR @@ -175,3 +177,163 @@ def test_select_usd_variants(): # Check if the variant selection is correct assert variant_set.GetVariantSelection() == "red" + + +def test_resolve_prim_pose(): + """Test resolve_prim_pose() function.""" + # number of objects + num_objects = 20 + # sample random scales for x, y, z + rand_scales = np.random.uniform(0.5, 1.5, size=(num_objects, 3, 3)) + rand_widths = np.random.uniform(0.1, 10.0, size=(num_objects,)) + rand_widths[:] = 1.0 + # sample random positions + rand_positions = np.random.uniform(-100, 100, size=(num_objects, 3, 3)) + # sample random rotations + rand_quats = np.random.randn(num_objects, 3, 4) + rand_quats /= np.linalg.norm(rand_quats, axis=2, keepdims=True) + + # create objects + for i in range(num_objects): + # simple cubes + cube_prim = prim_utils.create_prim( + f"/World/Cubes/instance_{i:02d}", + "Cube", + translation=rand_positions[i, 0], + orientation=rand_quats[i, 0], + scale=rand_scales[i, 0], + attributes={"size": rand_widths[i]}, + ) + # xform hierarchy + xform_prim = prim_utils.create_prim( + f"/World/Xform/instance_{i:02d}", + "Xform", + translation=rand_positions[i, 1], + orientation=rand_quats[i, 1], + # scale=rand_scales[i, 1], # FIXME: Enabling this affects the test to fail. + ) + geometry_prim = prim_utils.create_prim( + f"/World/Xform/instance_{i:02d}/geometry", + "Sphere", + translation=rand_positions[i, 2], + orientation=rand_quats[i, 2], + scale=rand_scales[i, 2], + attributes={"radius": rand_widths[i]}, + ) + dummy_prim = prim_utils.create_prim( + f"/World/Xform/instance_{i:02d}/dummy", + "Sphere", + ) + + # cube prim w.r.t. world frame + pos, quat = sim_utils.resolve_prim_pose(cube_prim) + pos, quat = np.array(pos), np.array(quat) + quat = quat if np.sign(rand_quats[i, 0, 0]) == np.sign(quat[0]) else -quat + np.testing.assert_allclose(pos, rand_positions[i, 0], atol=1e-3) + np.testing.assert_allclose(quat, rand_quats[i, 0], atol=1e-3) + # xform prim w.r.t. world frame + pos, quat = sim_utils.resolve_prim_pose(xform_prim) + pos, quat = np.array(pos), np.array(quat) + quat = quat if np.sign(rand_quats[i, 1, 0]) == np.sign(quat[0]) else -quat + np.testing.assert_allclose(pos, rand_positions[i, 1], atol=1e-3) + np.testing.assert_allclose(quat, rand_quats[i, 1], atol=1e-3) + # dummy prim w.r.t. world frame + pos, quat = sim_utils.resolve_prim_pose(dummy_prim) + pos, quat = np.array(pos), np.array(quat) + quat = quat if np.sign(rand_quats[i, 1, 0]) == np.sign(quat[0]) else -quat + np.testing.assert_allclose(pos, rand_positions[i, 1], atol=1e-3) + np.testing.assert_allclose(quat, rand_quats[i, 1], atol=1e-3) + + # geometry prim w.r.t. xform prim + pos, quat = sim_utils.resolve_prim_pose(geometry_prim, ref_prim=xform_prim) + pos, quat = np.array(pos), np.array(quat) + quat = quat if np.sign(rand_quats[i, 2, 0]) == np.sign(quat[0]) else -quat + np.testing.assert_allclose(pos, rand_positions[i, 2], atol=1e-3) + np.testing.assert_allclose(quat, rand_quats[i, 2], atol=1e-3) + # dummy prim w.r.t. xform prim + pos, quat = sim_utils.resolve_prim_pose(dummy_prim, ref_prim=xform_prim) + pos, quat = np.array(pos), np.array(quat) + np.testing.assert_allclose(pos, np.zeros(3), atol=1e-3) + np.testing.assert_allclose(quat, np.array([1, 0, 0, 0]), atol=1e-3) + # xform prim w.r.t. cube prim + pos, quat = sim_utils.resolve_prim_pose(xform_prim, ref_prim=cube_prim) + pos, quat = np.array(pos), np.array(quat) + # -- compute ground truth values + gt_pos, gt_quat = math_utils.subtract_frame_transforms( + torch.from_numpy(rand_positions[i, 0]).unsqueeze(0), + torch.from_numpy(rand_quats[i, 0]).unsqueeze(0), + torch.from_numpy(rand_positions[i, 1]).unsqueeze(0), + torch.from_numpy(rand_quats[i, 1]).unsqueeze(0), + ) + gt_pos, gt_quat = gt_pos.squeeze(0).numpy(), gt_quat.squeeze(0).numpy() + quat = quat if np.sign(gt_quat[0]) == np.sign(quat[0]) else -quat + np.testing.assert_allclose(pos, gt_pos, atol=1e-3) + np.testing.assert_allclose(quat, gt_quat, atol=1e-3) + + +def test_resolve_prim_scale(): + """Test resolve_prim_scale() function. + + To simplify the test, we assume that the effective scale at a prim + is the product of the scales of the prims in the hierarchy: + + scale = scale_of_xform * scale_of_geometry_prim + + This is only true when rotations are identity or the transforms are + orthogonal and uniformly scaled. Otherwise, scale is not composable + like that in local component-wise fashion. + """ + # number of objects + num_objects = 20 + # sample random scales for x, y, z + rand_scales = np.random.uniform(0.5, 1.5, size=(num_objects, 3, 3)) + rand_widths = np.random.uniform(0.1, 10.0, size=(num_objects,)) + rand_widths[:] = 1.0 + # sample random positions + rand_positions = np.random.uniform(-100, 100, size=(num_objects, 3, 3)) + + # create objects + for i in range(num_objects): + # simple cubes + cube_prim = prim_utils.create_prim( + f"/World/Cubes/instance_{i:02d}", + "Cube", + translation=rand_positions[i, 0], + scale=rand_scales[i, 0], + attributes={"size": rand_widths[i]}, + ) + # xform hierarchy + xform_prim = prim_utils.create_prim( + f"/World/Xform/instance_{i:02d}", + "Xform", + translation=rand_positions[i, 1], + scale=rand_scales[i, 1], + ) + geometry_prim = prim_utils.create_prim( + f"/World/Xform/instance_{i:02d}/geometry", + "Sphere", + translation=rand_positions[i, 2], + scale=rand_scales[i, 2], + attributes={"radius": rand_widths[i]}, + ) + dummy_prim = prim_utils.create_prim( + f"/World/Xform/instance_{i:02d}/dummy", + "Sphere", + ) + + # cube prim + scale = sim_utils.resolve_prim_scale(cube_prim) + scale = np.array(scale) + np.testing.assert_allclose(scale, rand_scales[i, 0], atol=1e-5) + # xform prim + scale = sim_utils.resolve_prim_scale(xform_prim) + scale = np.array(scale) + np.testing.assert_allclose(scale, rand_scales[i, 1], atol=1e-5) + # geometry prim + scale = sim_utils.resolve_prim_scale(geometry_prim) + scale = np.array(scale) + np.testing.assert_allclose(scale, rand_scales[i, 1] * rand_scales[i, 2], atol=1e-5) + # dummy prim + scale = sim_utils.resolve_prim_scale(dummy_prim) + scale = np.array(scale) + np.testing.assert_allclose(scale, rand_scales[i, 1], atol=1e-5) diff --git a/source/isaaclab_rl/isaaclab_rl/rsl_rl/rl_cfg.py b/source/isaaclab_rl/isaaclab_rl/rsl_rl/rl_cfg.py index 5b03a7c639b..87ce3037fef 100644 --- a/source/isaaclab_rl/isaaclab_rl/rsl_rl/rl_cfg.py +++ b/source/isaaclab_rl/isaaclab_rl/rsl_rl/rl_cfg.py @@ -64,6 +64,43 @@ class RslRlPpoActorCriticRecurrentCfg(RslRlPpoActorCriticCfg): """The number of RNN layers.""" +@configclass +class RslRlPerceptiveActorCriticCfg(RslRlPpoActorCriticCfg): + """Configuration for the PPO actor-critic networks with perceptual layers.""" + + @configclass + class CNNConfig: + out_channels: list[int] = MISSING + """The number of output channels for the CNN.""" + + kernel_size: list[tuple[int, int]] | tuple[int, int] = MISSING + """The kernel size for the CNN.""" + + stride: list[int] | int = 1 + """The stride for the CNN.""" + + flatten: bool = True + """Whether to flatten the output of the CNN.""" + + avg_pool: tuple[int, int] | None = None + """The average pool for the CNN.""" + + batchnorm: bool | list[bool] = False + """Whether to use batch normalization for the CNN.""" + + max_pool: bool | list[bool] = False + """Whether to use max pooling for the CNN.""" + + class_name: str = "PerceptiveActorCritic" + """The policy class name. Default is PerceptiveActorCritic.""" + + actor_cnn_config: list[CNNConfig] | CNNConfig | None = MISSING + """The CNN configuration for the actor network.""" + + critic_cnn_config: list[CNNConfig] | CNNConfig | None = MISSING + """The CNN configuration for the critic network.""" + + ############################ # Algorithm configurations # ############################ diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/config/anymal_c/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/config/anymal_c/__init__.py index 84a6dd920c3..9f902d1abba 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/config/anymal_c/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/config/anymal_c/__init__.py @@ -12,23 +12,41 @@ ## gym.register( - id="Isaac-Navigation-Flat-Anymal-C-v0", + id="Isaac-Navigation-Anymal-C-RayCaster-v0", entry_point="isaaclab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": f"{__name__}.navigation_env_cfg:NavigationEnvCfg", + "env_cfg_entry_point": f"{__name__}.env_cfg:AnymalCRayCasterNavEnvCfg", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:NavigationEnvPPORunnerCfg", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", }, ) gym.register( - id="Isaac-Navigation-Flat-Anymal-C-Play-v0", + id="Isaac-Navigation-Anymal-C-RayCaster-Play-v0", entry_point="isaaclab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": f"{__name__}.navigation_env_cfg:NavigationEnvCfg_PLAY", + "env_cfg_entry_point": f"{__name__}.env_cfg:AnymalCRayCasterNavEnvCfg_PLAY", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:NavigationEnvPPORunnerCfg", + }, +) + +gym.register( + id="Isaac-Navigation-Anymal-C-Tiled-v0", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.env_cfg:AnymalCTiledNavEnvCfg", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:NavigationEnvPPORunnerCfg", + }, +) + +gym.register( + id="Isaac-Navigation-Anymal-C-Tiled-Play-v0", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.env_cfg:AnymalCTiledNavEnvCfg_PLAY", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:NavigationEnvPPORunnerCfg", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", }, ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/config/anymal_c/agents/rsl_rl_ppo_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/config/anymal_c/agents/rsl_rl_ppo_cfg.py index 4b23def89b2..ea258b7f68f 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/config/anymal_c/agents/rsl_rl_ppo_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/config/anymal_c/agents/rsl_rl_ppo_cfg.py @@ -5,34 +5,11 @@ from isaaclab.utils import configclass -from isaaclab_rl.rsl_rl import RslRlOnPolicyRunnerCfg, RslRlPpoActorCriticCfg, RslRlPpoAlgorithmCfg +from isaaclab_tasks.manager_based.navigation.config.base_rsl_rl_ppo_cfg import NavBasePPORunnerCfg @configclass -class NavigationEnvPPORunnerCfg(RslRlOnPolicyRunnerCfg): - num_steps_per_env = 8 +class NavigationEnvPPORunnerCfg(NavBasePPORunnerCfg): max_iterations = 1500 save_interval = 50 experiment_name = "anymal_c_navigation" - policy = RslRlPpoActorCriticCfg( - init_noise_std=0.5, - actor_obs_normalization=False, - critic_obs_normalization=False, - actor_hidden_dims=[128, 128], - critic_hidden_dims=[128, 128], - activation="elu", - ) - algorithm = RslRlPpoAlgorithmCfg( - value_loss_coef=1.0, - use_clipped_value_loss=True, - clip_param=0.2, - entropy_coef=0.005, - num_learning_epochs=5, - num_mini_batches=4, - learning_rate=1.0e-3, - schedule="adaptive", - gamma=0.99, - lam=0.95, - desired_kl=0.01, - max_grad_norm=1.0, - ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/config/anymal_c/agents/skrl_flat_ppo_cfg.yaml b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/config/anymal_c/agents/skrl_flat_ppo_cfg.yaml deleted file mode 100644 index 005f95806d1..00000000000 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/config/anymal_c/agents/skrl_flat_ppo_cfg.yaml +++ /dev/null @@ -1,85 +0,0 @@ -# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -seed: 42 - - -# Models are instantiated using skrl's model instantiator utility -# https://skrl.readthedocs.io/en/latest/api/utils/model_instantiators.html -models: - separate: False - policy: # see gaussian_model parameters - class: GaussianMixin - clip_actions: False - clip_log_std: True - min_log_std: -20.0 - max_log_std: 2.0 - initial_log_std: -0.6931471805599453 - network: - - name: net - input: STATES - layers: [128, 128] - activations: elu - output: ACTIONS - value: # see deterministic_model parameters - class: DeterministicMixin - clip_actions: False - network: - - name: net - input: STATES - layers: [128, 128] - activations: elu - output: ONE - - -# Rollout memory -# https://skrl.readthedocs.io/en/latest/api/memories/random.html -memory: - class: RandomMemory - memory_size: -1 # automatically determined (same as agent:rollouts) - - -# PPO agent configuration (field names are from PPO_DEFAULT_CONFIG) -# https://skrl.readthedocs.io/en/latest/api/agents/ppo.html -agent: - class: PPO - rollouts: 8 - learning_epochs: 5 - mini_batches: 4 - discount_factor: 0.99 - lambda: 0.95 - learning_rate: 1.0e-03 - learning_rate_scheduler: KLAdaptiveLR - learning_rate_scheduler_kwargs: - kl_threshold: 0.01 - state_preprocessor: null - state_preprocessor_kwargs: null - value_preprocessor: null - value_preprocessor_kwargs: null - random_timesteps: 0 - learning_starts: 0 - grad_norm_clip: 1.0 - ratio_clip: 0.2 - value_clip: 0.2 - clip_predicted_values: True - entropy_loss_scale: 0.005 - value_loss_scale: 1.0 - kl_threshold: 0.0 - rewards_shaper_scale: 1.0 - time_limit_bootstrap: False - # logging and checkpoint - experiment: - directory: "anymal_c_navigation" - experiment_name: "" - write_interval: auto - checkpoint_interval: auto - - -# Sequential trainer -# https://skrl.readthedocs.io/en/latest/api/trainers/sequential.html -trainer: - class: SequentialTrainer - timesteps: 12000 - environment_info: log diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/config/anymal_c/env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/config/anymal_c/env_cfg.py new file mode 100644 index 00000000000..aa61533e97e --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/config/anymal_c/env_cfg.py @@ -0,0 +1,69 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from isaaclab.utils import configclass + +from isaaclab_tasks.manager_based.navigation.navigation_env_cfg import RayCasterNavEnvCfg, TiledNavEnvCfg + +## +# Pre-defined configs +## +from isaaclab_assets.robots.anymal import ANYMAL_C_CFG # isort: skip + + +def _play_init(self): + # make a smaller scene for play + self.scene.num_envs = 50 + self.scene.env_spacing = 2.5 + # spawn the robot randomly in the grid (instead of their terrain levels) + self.scene.terrain.max_init_terrain_level = None + # reduce the number of terrains to save memory + if self.scene.terrain.terrain_generator is not None: + self.scene.terrain.terrain_generator.num_rows = 5 + self.scene.terrain.terrain_generator.num_cols = 5 + self.scene.terrain.terrain_generator.curriculum = False + + # disable randomization for play + self.observations.policy.enable_corruption = False + + +@configclass +class AnymalCRayCasterNavEnvCfg(RayCasterNavEnvCfg): + def __post_init__(self): + # post init of parent + super().__post_init__() + # switch robot to anymal-c + self.scene.robot = ANYMAL_C_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + # turn off the self-collisions + self.scene.robot.spawn.articulation_props.enabled_self_collisions = False + + +@configclass +class AnymalCTiledNavEnvCfg(TiledNavEnvCfg): + def __post_init__(self): + # post init of parent + super().__post_init__() + # switch robot to anymal-c + self.scene.robot = ANYMAL_C_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + # turn off the self-collisions + self.scene.robot.spawn.articulation_props.enabled_self_collisions = False + + +@configclass +class AnymalCRayCasterNavEnvCfg_PLAY(AnymalCRayCasterNavEnvCfg): + def __post_init__(self): + # post init of parent + super().__post_init__() + + _play_init(self) + + +@configclass +class AnymalCTiledNavEnvCfg_PLAY(AnymalCTiledNavEnvCfg): + def __post_init__(self): + # post init of parent + super().__post_init__() + + _play_init(self) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/config/anymal_c/navigation_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/config/anymal_c/navigation_env_cfg.py deleted file mode 100644 index 2f2162fdce9..00000000000 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/config/anymal_c/navigation_env_cfg.py +++ /dev/null @@ -1,160 +0,0 @@ -# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -import math - -from isaaclab.envs import ManagerBasedRLEnvCfg -from isaaclab.managers import EventTermCfg as EventTerm -from isaaclab.managers import ObservationGroupCfg as ObsGroup -from isaaclab.managers import ObservationTermCfg as ObsTerm -from isaaclab.managers import RewardTermCfg as RewTerm -from isaaclab.managers import SceneEntityCfg -from isaaclab.managers import TerminationTermCfg as DoneTerm -from isaaclab.utils import configclass -from isaaclab.utils.assets import ISAACLAB_NUCLEUS_DIR - -import isaaclab_tasks.manager_based.navigation.mdp as mdp -from isaaclab_tasks.manager_based.locomotion.velocity.config.anymal_c.flat_env_cfg import AnymalCFlatEnvCfg - -LOW_LEVEL_ENV_CFG = AnymalCFlatEnvCfg() - - -@configclass -class EventCfg: - """Configuration for events.""" - - reset_base = EventTerm( - func=mdp.reset_root_state_uniform, - mode="reset", - params={ - "pose_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5), "yaw": (-3.14, 3.14)}, - "velocity_range": { - "x": (-0.0, 0.0), - "y": (-0.0, 0.0), - "z": (-0.0, 0.0), - "roll": (-0.0, 0.0), - "pitch": (-0.0, 0.0), - "yaw": (-0.0, 0.0), - }, - }, - ) - - -@configclass -class ActionsCfg: - """Action terms for the MDP.""" - - pre_trained_policy_action: mdp.PreTrainedPolicyActionCfg = mdp.PreTrainedPolicyActionCfg( - asset_name="robot", - policy_path=f"{ISAACLAB_NUCLEUS_DIR}/Policies/ANYmal-C/Blind/policy.pt", - low_level_decimation=4, - low_level_actions=LOW_LEVEL_ENV_CFG.actions.joint_pos, - low_level_observations=LOW_LEVEL_ENV_CFG.observations.policy, - ) - - -@configclass -class ObservationsCfg: - """Observation specifications for the MDP.""" - - @configclass - class PolicyCfg(ObsGroup): - """Observations for policy group.""" - - # observation terms (order preserved) - base_lin_vel = ObsTerm(func=mdp.base_lin_vel) - projected_gravity = ObsTerm(func=mdp.projected_gravity) - pose_command = ObsTerm(func=mdp.generated_commands, params={"command_name": "pose_command"}) - - # observation groups - policy: PolicyCfg = PolicyCfg() - - -@configclass -class RewardsCfg: - """Reward terms for the MDP.""" - - termination_penalty = RewTerm(func=mdp.is_terminated, weight=-400.0) - position_tracking = RewTerm( - func=mdp.position_command_error_tanh, - weight=0.5, - params={"std": 2.0, "command_name": "pose_command"}, - ) - position_tracking_fine_grained = RewTerm( - func=mdp.position_command_error_tanh, - weight=0.5, - params={"std": 0.2, "command_name": "pose_command"}, - ) - orientation_tracking = RewTerm( - func=mdp.heading_command_error_abs, - weight=-0.2, - params={"command_name": "pose_command"}, - ) - - -@configclass -class CommandsCfg: - """Command terms for the MDP.""" - - pose_command = mdp.UniformPose2dCommandCfg( - asset_name="robot", - simple_heading=False, - resampling_time_range=(8.0, 8.0), - debug_vis=True, - ranges=mdp.UniformPose2dCommandCfg.Ranges(pos_x=(-3.0, 3.0), pos_y=(-3.0, 3.0), heading=(-math.pi, math.pi)), - ) - - -@configclass -class TerminationsCfg: - """Termination terms for the MDP.""" - - time_out = DoneTerm(func=mdp.time_out, time_out=True) - base_contact = DoneTerm( - func=mdp.illegal_contact, - params={"sensor_cfg": SceneEntityCfg("contact_forces", body_names="base"), "threshold": 1.0}, - ) - - -@configclass -class NavigationEnvCfg(ManagerBasedRLEnvCfg): - """Configuration for the navigation environment.""" - - # environment settings - scene: SceneEntityCfg = LOW_LEVEL_ENV_CFG.scene - actions: ActionsCfg = ActionsCfg() - observations: ObservationsCfg = ObservationsCfg() - events: EventCfg = EventCfg() - # mdp settings - commands: CommandsCfg = CommandsCfg() - rewards: RewardsCfg = RewardsCfg() - terminations: TerminationsCfg = TerminationsCfg() - - def __post_init__(self): - """Post initialization.""" - - self.sim.dt = LOW_LEVEL_ENV_CFG.sim.dt - self.sim.render_interval = LOW_LEVEL_ENV_CFG.decimation - self.decimation = LOW_LEVEL_ENV_CFG.decimation * 10 - self.episode_length_s = self.commands.pose_command.resampling_time_range[1] - - if self.scene.height_scanner is not None: - self.scene.height_scanner.update_period = ( - self.actions.pre_trained_policy_action.low_level_decimation * self.sim.dt - ) - if self.scene.contact_forces is not None: - self.scene.contact_forces.update_period = self.sim.dt - - -class NavigationEnvCfg_PLAY(NavigationEnvCfg): - def __post_init__(self) -> None: - # post init of parent - super().__post_init__() - - # make a smaller scene for play - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - # disable randomization for play - self.observations.policy.enable_corruption = False diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/config/base_rsl_rl_ppo_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/config/base_rsl_rl_ppo_cfg.py new file mode 100644 index 00000000000..eea79eefebf --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/config/base_rsl_rl_ppo_cfg.py @@ -0,0 +1,56 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from isaaclab.utils import configclass + +from isaaclab_rl.rsl_rl import RslRlOnPolicyRunnerCfg, RslRlPerceptiveActorCriticCfg, RslRlPpoAlgorithmCfg + + +@configclass +class NavBasePPORunnerCfg(RslRlOnPolicyRunnerCfg): + num_steps_per_env = 24 + max_iterations = 1500 + save_interval = 50 + experiment_name = "nav_base" + obs_groups = { + "policy": ["proprioceptive", "exteroceptive"], + "critic": ["proprioceptive", "exteroceptive"], + } + policy = RslRlPerceptiveActorCriticCfg( + init_noise_std=1.0, + actor_obs_normalization=False, + critic_obs_normalization=False, + actor_hidden_dims=[512, 256, 128], + critic_hidden_dims=[512, 256, 128], + actor_cnn_config=RslRlPerceptiveActorCriticCfg.CNNConfig( + out_channels=[32, 64], + kernel_size=[(7, 7), (5, 5)], + flatten=False, + avg_pool=(1, 1), + max_pool=(True, False), + ), + critic_cnn_config=RslRlPerceptiveActorCriticCfg.CNNConfig( + out_channels=[32, 64], + kernel_size=[(7, 7), (5, 5)], + flatten=False, + avg_pool=(1, 1), + max_pool=(True, False), + ), + activation="elu", + ) + algorithm = RslRlPpoAlgorithmCfg( + value_loss_coef=1.0, + use_clipped_value_loss=True, + clip_param=0.2, + entropy_coef=0.005, + num_learning_epochs=5, + num_mini_batches=4, + learning_rate=1.0e-3, + schedule="adaptive", + gamma=0.99, + lam=0.95, + desired_kl=0.01, + max_grad_norm=1.0, + ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/__init__.py index 0a2576ceb9f..e0713d3acc4 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/__init__.py @@ -7,5 +7,9 @@ from isaaclab.envs.mdp import * # noqa: F401, F403 -from .pre_trained_policy_action import * # noqa: F401, F403 +from .actions import * # noqa: F401, F403 +from .commands import * # noqa: F401, F403 +from .events import * # noqa: F401, F403 +from .observations import * # noqa: F401, F403 from .rewards import * # noqa: F401, F403 +from .terminations import * # noqa: F401, F403 diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/actions/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/actions/__init__.py new file mode 100644 index 00000000000..a44f57710b8 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/actions/__init__.py @@ -0,0 +1,12 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Copyright (c) 2025, The Nav-Suite Project Developers (https://github.com/leggedrobotics/nav-suite/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from .navigation_actions import NavigationSE2Action +from .navigation_actions_cfg import NavigationSE2ActionCfg diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/actions/navigation_actions.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/actions/navigation_actions.py new file mode 100644 index 00000000000..d4dc2f8c5fe --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/actions/navigation_actions.py @@ -0,0 +1,181 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Code adapted from https://github.com/leggedrobotics/nav-suite + +# Copyright (c) 2025, The Nav-Suite Project Developers (https://github.com/leggedrobotics/nav-suite/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import torch +from typing import TYPE_CHECKING + +from isaaclab.envs import ManagerBasedRLEnv +from isaaclab.managers.action_manager import ActionTerm +from isaaclab.utils.assets import check_file_path, read_file + +if TYPE_CHECKING: + from .navigation_actions_cfg import NavigationSE2ActionCfg + + +class NavigationSE2Action(ActionTerm): + """Actions to navigate a robot by following some path.""" + + cfg: NavigationSE2ActionCfg + _env: ManagerBasedRLEnv + + def __init__(self, cfg: NavigationSE2ActionCfg, env: ManagerBasedRLEnv): + super().__init__(cfg, env) + # check if policy file exists + if not check_file_path(cfg.low_level_policy_file): + raise FileNotFoundError(f"Policy file '{cfg.low_level_policy_file}' does not exist.") + # load policies + file_bytes = read_file(self.cfg.low_level_policy_file) + self.low_level_policy = torch.jit.load(file_bytes, map_location=self.device) + self.low_level_policy.eval() + if self.cfg.freeze_low_level_policy: + self.low_level_policy = torch.jit.freeze(self.low_level_policy) + + # prepare joint position actions + if not isinstance(self.cfg.low_level_action, list): + self.cfg.low_level_action = [self.cfg.low_level_action] + self.low_level_action_terms: list[ActionTerm] = [ + term_cfg.class_type(term_cfg, env) for term_cfg in self.cfg.low_level_action + ] + + # parse scale + if self.cfg.scale is not None: + if isinstance(self.cfg.scale, (float, int)): + self._scale = float(self.cfg.scale) + elif isinstance(self.cfg.scale, list): + self._scale = torch.tensor([self.cfg.scale], device=self.device).repeat(self.num_envs, 1) + else: + raise ValueError( + f"Unsupported scale type: {type(self.cfg.scale)}. Supported types are float, int, and list." + ) + else: + self._scale = 1 + + # parse offset + if self.cfg.offset is not None: + if isinstance(self.cfg.offset, (float, int)): + self._offset = float(self.cfg.offset) + elif isinstance(self.cfg.offset, list): + self._offset = torch.tensor([self.cfg.offset], device=self.device).repeat(self.num_envs, 1) + else: + raise ValueError( + f"Unsupported offset type: {type(self.cfg.offset)}. Supported types are float, int, and list." + ) + else: + self._offset = 0 + + # parse clip + if self.cfg.clip_mode != "none": + if self.cfg.clip_mode == "minmax": + assert isinstance( + self.cfg.clip, list + ), "Clip must be a list of tuples of (min, max) values if clip mode is 'minmax'" + assert ( + len(self.cfg.clip) == self.cfg.action_dim + ), "Clip must have the same length as the action dimension" + self._clip = torch.tensor(self.cfg.clip, device=self.device).repeat(self.num_envs, 1, 1) + elif self.cfg.clip_mode == "tanh": + pass + else: + raise ValueError( + f"Unsupported clip mode: {self.cfg.clip_mode}. Supported modes are 'minmax', 'tanh' and 'none'." + ) + + # set up buffers + self._init_buffers() + + """ + Properties. + """ + + @property + def action_dim(self) -> int: + return self.cfg.action_dim + + @property + def raw_actions(self) -> torch.Tensor: + return self._raw_navigation_velocity_actions + + @property + def processed_actions(self) -> torch.Tensor: + return self._processed_navigation_velocity_actions + + @property + def low_level_actions(self) -> torch.Tensor: + return self._low_level_actions + + @property + def prev_low_level_actions(self) -> torch.Tensor: + return self._prev_low_level_actions + + """ + Operations. + """ + + def process_actions(self, actions): + """Process low-level navigation actions. This function is called with a frequency of 10Hz""" + + # Store the navigation actions + self._raw_navigation_velocity_actions[:] = actions + self._processed_navigation_velocity_actions[:] = actions.clone().view(self.num_envs, self.cfg.action_dim) + # clip actions + if self.cfg.clip_mode == "minmax": + self._processed_navigation_velocity_actions = torch.clamp( + self._processed_navigation_velocity_actions, min=self._clip[:, :, 0], max=self._clip[:, :, 1] + ) + elif self.cfg.clip_mode == "tanh": + self._processed_navigation_velocity_actions = torch.tanh(self._processed_navigation_velocity_actions) + # apply the affine transformations + self._processed_navigation_velocity_actions *= self._scale + self._processed_navigation_velocity_actions += self._offset + + def apply_actions(self): + """Apply low-level actions for the simulator to the physics engine. This functions is called with the + simulation frequency of 200Hz. Since low-level locomotion runs at 50Hz, we need to decimate the actions.""" + + if self._counter % self.cfg.low_level_decimation == 0: + self._counter = 0 + self._prev_low_level_actions[:] = self._low_level_actions.clone() + # Get low level actions from low level policy + self._low_level_actions[:] = self.low_level_policy( + self._env.observation_manager.compute_group(group_name=self.cfg.low_level_obs_group) + ) + + # split the actions and apply to each tensor + idx = 0 + for term in self.low_level_action_terms: + term_actions = self._low_level_actions[:, idx : idx + term.action_dim] + term.process_actions(term_actions) + idx += term.action_dim + + # Apply low level actions + for term in self.low_level_action_terms: + term.apply_actions() + self._counter += 1 + + """ + Helper functions + """ + + def _init_buffers(self): + # Prepare buffers + self._raw_navigation_velocity_actions = torch.zeros(self.num_envs, self.cfg.action_dim, device=self.device) + self._processed_navigation_velocity_actions = torch.zeros( + (self.num_envs, self.cfg.action_dim), device=self.device + ) + self._low_level_actions = torch.zeros( + self.num_envs, sum([term.action_dim for term in self.low_level_action_terms]), device=self.device + ) + self._prev_low_level_actions = torch.zeros_like(self._low_level_actions) + self._low_level_step_dt = self.cfg.low_level_decimation * self._env.physics_dt + self._counter = 0 diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/actions/navigation_actions_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/actions/navigation_actions_cfg.py new file mode 100644 index 00000000000..549b1240e57 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/actions/navigation_actions_cfg.py @@ -0,0 +1,74 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Code adapted from https://github.com/leggedrobotics/nav-suite + +# Copyright (c) 2025, The Nav-Suite Project Developers (https://github.com/leggedrobotics/nav-suite/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from dataclasses import MISSING +from typing import Literal + +from isaaclab.managers.action_manager import ActionTerm, ActionTermCfg +from isaaclab.utils import configclass + +from .navigation_actions import NavigationSE2Action + + +@configclass +class NavigationSE2ActionCfg(ActionTermCfg): + class_type: type[ActionTerm] = NavigationSE2Action + """ Class of the action term.""" + + low_level_decimation: int = 4 + """Decimation factor for the low level action term.""" + + low_level_action: ActionTermCfg | list[ActionTermCfg] = MISSING + """Configuration of the low level action term.""" + + low_level_policy_file: str = MISSING + """Path to the low level policy file.""" + + freeze_low_level_policy: bool = True + """Whether to freeze the low level policy. + + Can improve performance but will also eliminate possible functions such as `reset`.""" + + low_level_obs_group: str = "low_level_policy" + """Observation group of the low level policy.""" + + action_dim: int = 3 + """Dimension of the action space. Default is 3 [vx, vy, omega].""" + + clip_mode: Literal["none", "minmax", "tanh"] = "none" + """Clip mode for the action space. Default is "none". + + "minmax": Clip the action space to the range of the min and max values defined in :attr:`clip`. + "tanh": Clip the action space to the range of the tanh function, i.e., each action dimension between -1 and 1. + "none": No clipping is applied. + """ + + clip: list[tuple[float, float]] | None = None + """Clip the action space. Default is None. + + For each action dimension, provide a tuple of (min, max) values.""" + + scale: float | list[float] | None = None + """Scale the action space. Default is None. + + .. note:: + Scale is applied after clipping. If a list is provided, it must be of the same length as the number of action dimensions. + """ + + offset: float | list[float] | None = None + """Offset the action space. Default is None. + + .. note:: + Offset is applied after scaling. If a list is provided, it must be of the same length as the number of action dimensions. + """ diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/commands/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/commands/__init__.py new file mode 100644 index 00000000000..b14e9664751 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/commands/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from .goal_command import GoalCommandTerm +from .goal_command_cfg import GoalCommandCfg diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/commands/goal_command.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/commands/goal_command.py new file mode 100644 index 00000000000..254bb12b498 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/commands/goal_command.py @@ -0,0 +1,372 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Code adapted from https://github.com/leggedrobotics/nav-suite + +# Copyright (c) 2025, The Nav-Suite Project Developers (https://github.com/leggedrobotics/nav-suite/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + + +from __future__ import annotations + +import numpy as np +import torch +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import isaaclab.sim as sim_utils +from isaaclab.assets import Articulation +from isaaclab.managers import CommandTerm +from isaaclab.markers import VisualizationMarkers +from isaaclab.markers.config import CUBOID_MARKER_CFG +from isaaclab.markers.visualization_markers import VisualizationMarkersCfg +from isaaclab.sensors import RayCaster, RayCasterCamera +from isaaclab.utils.math import quat_apply_inverse, quat_from_angle_axis, wrap_to_pi, yaw_quat +from isaaclab.utils.warp import raycast_mesh + +CYLINDER_MARKER_CFG = VisualizationMarkersCfg( + markers={ + "cylinder": sim_utils.CylinderCfg( + radius=1, + height=1, + axis="X", + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.0, 1.0, 0.0)), + ), + } +) + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedRLEnv + + from .goal_command_cfg import GoalCommandCfg + +# TODO: @pascal-roth: remove once multi-mesh raycasting is merged +try: + from isaaclab.utils.warp import raycast_dynamic_meshes +except ImportError: + raycast_dynamic_meshes = None + + +class GoalCommandTerm(CommandTerm): + r"""Base class for goal commands. + + This class is used to define the common visualization features for goal commands. + """ + + cfg: GoalCommandCfg + """Configuration for the command.""" + + def __init__(self, cfg: GoalCommandCfg, env: ManagerBasedRLEnv): + super().__init__(cfg, env) + + # -- robot + self.robot: Articulation = env.scene[cfg.asset_name] + + # -- generate height map + self._construct_height_map() + # -- construct traversability map + self._construct_traversability_map() + + # -- goal commands: (x, y, z) + self.pos_spawn_w = torch.zeros(self.num_envs, 3, device=self.device) + self.heading_spawn_w = torch.zeros(self.num_envs, device=self.device) + self.pos_command_w = torch.zeros(self.num_envs, 3, device=self.device) + self.heading_command_w = torch.zeros(self.num_envs, device=self.device) + self.pos_command_b = torch.zeros_like(self.pos_command_w) + self.heading_command_b = torch.zeros_like(self.heading_command_w) + + def __str__(self) -> str: + msg = "GoalCommandGenerator:\n" + msg += f"\tCommand dimension:\t {tuple(self.command.shape[1:])}\n" + msg += f"\tResampling time range:\t {self.cfg.resampling_time_range}\n" + return msg + + """ + Properties + """ + + @property + def command(self) -> torch.Tensor: + """The desired base pose in base frame. Shape is (num_envs, 4).""" + return torch.cat((self.pos_command_b, self.heading_command_b.unsqueeze(-1)), dim=1) + + """ + Implementation specific functions. + """ + + def _resample_command(self, env_ids: Sequence[int]): + """Resample the command for the specified environments.""" + + # get the terrain id for the environments + terrain_id = ( + torch.cdist(self._env.scene.terrain.env_origins[env_ids], self._terrain_origins, p=2) + .argmin(dim=1) + .unsqueeze(0) + ) + + start_sample = torch.concat( + (terrain_id, torch.randint(0, self.split_max_length, (len(env_ids),), device=self.device).unsqueeze(0)) + ) + end_sample = torch.concat( + (terrain_id, torch.randint(0, self.split_max_length, (len(env_ids),), device=self.device).unsqueeze(0)) + ) + + # robot height + robot_height = self.robot.data.default_root_state[env_ids, 2] + + # Update command buffers + self.pos_command_w[env_ids] = self._split_traversability_map[start_sample[0, :], start_sample[1, :]] + self.pos_command_w[env_ids, 2] = robot_height + + # Update spawn locations and heading buffer + self.pos_spawn_w[env_ids] = self._split_traversability_map[end_sample[0, :], end_sample[1, :]] + self.pos_spawn_w[env_ids, 2] = robot_height + + # Calculate the spawn heading based on the goal position + self.heading_spawn_w[env_ids] = torch.atan2( + self.pos_command_w[env_ids, 1] - self.pos_spawn_w[env_ids, 1], + self.pos_command_w[env_ids, 0] - self.pos_spawn_w[env_ids, 0], + ) + # Calculate the goal heading based on the goal position + self.heading_command_w[env_ids] = torch.atan2( + self.pos_command_w[env_ids, 1] - self.pos_spawn_w[env_ids, 1], + self.pos_command_w[env_ids, 0] - self.pos_spawn_w[env_ids, 0], + ) + + # NOTE: the reset event is called before the new goal commands are generated, i.e. the spawn locations are + # updated before the new goal commands are generated. To repsawn with the correct locations, we call here the + # update spawn locations function + if self.cfg.reset_pos_term_name: + reset_term_idx = self._env.event_manager.active_terms["reset"].index(self.cfg.reset_pos_term_name) + self._env.event_manager._mode_term_cfgs["reset"][reset_term_idx].func( + self._env, env_ids, **self._env.event_manager._mode_term_cfgs["reset"][reset_term_idx].params + ) + + def _update_command(self): + """Re-target the position command to the current root position and heading.""" + target_vec = self.pos_command_w - self.robot.data.root_pos_w[:, :3] + target_vec[:, 2] = 0.0 # ignore z component + self.pos_command_b[:] = quat_apply_inverse(yaw_quat(self.robot.data.root_quat_w), target_vec) + + # update the heading command in the base frame + # heading_w is angle world x axis to robot base x axis + self.heading_command_b[:] = wrap_to_pi(self.heading_command_w - self.robot.data.heading_w) + + def _update_metrics(self): + """Update metrics.""" + self.metrics["error_pos"] = torch.norm(self.pos_command_w - self.robot.data.root_pos_w[:, :3], dim=1) + + """ + Helper functions + """ + + def _get_mesh_dimensions( + self, raycaster: RayCaster | RayCasterCamera + ) -> tuple[float, float, float, float, float, float]: + # get min, max of the mesh in the xy plane + # Get bounds of the terrain + bounds = [] + + def deep_flatten(nested_list): + for item in nested_list: + if isinstance(item, list): + yield from deep_flatten(item) # recurse into sublist + else: + yield item + + for mesh in list(deep_flatten(raycaster.meshes.values())): + curr_bounds = torch.zeros((2, 3)) + curr_bounds[0] = torch.tensor(mesh.points).max(dim=0)[0] + curr_bounds[1] = torch.tensor(mesh.points).min(dim=0)[0] + bounds.append(curr_bounds) + bounds = torch.vstack(bounds) + x_min, y_min, z_min = bounds[:, 0].min().item(), bounds[:, 1].min().item(), bounds[:, 2].min().item() + x_max, y_max, z_max = bounds[:, 0].max().item(), bounds[:, 1].max().item(), bounds[:, 2].max().item() + return x_max, y_max, x_min, y_min, z_min, z_max + + def _construct_height_map(self): + # get dimensions and construct height grid with raycasting + raycaster: RayCaster | RayCasterCamera = self._env.scene.sensors[self.cfg.raycaster_sensor] + + # get mesh dimensions [x_max, y_max, x_min, y_min] + mesh_dimensions = self._get_mesh_dimensions(raycaster) + + grid_x, grid_y = torch.meshgrid( + torch.linspace( + mesh_dimensions[2], + mesh_dimensions[0], + int(np.abs(np.ceil((mesh_dimensions[0] - mesh_dimensions[2]) / self.cfg.grid_resolution))), + device=self.device, + ), + torch.linspace( + mesh_dimensions[3], + mesh_dimensions[1], + int(np.abs(np.ceil((mesh_dimensions[1] - mesh_dimensions[3]) / self.cfg.grid_resolution))), + device=self.device, + ), + indexing="ij", + ) + grid_z = torch.ones_like(grid_x) * raycaster.cfg.max_distance + self._height_grid_pos = torch.vstack((grid_x.flatten(), grid_y.flatten(), grid_z.flatten())).T + direction = torch.zeros_like(self._height_grid_pos) + direction[:, 2] = -1.0 + + # check for collision with raycasting from the top + # support for both multi-mesh and single-mesh raycasting + if hasattr(raycaster, "_mesh_ids_wp") and raycast_dynamic_meshes is not None: + hit_point = raycast_dynamic_meshes( + ray_starts=self._height_grid_pos.unsqueeze(0), + ray_directions=direction.unsqueeze(0), + max_dist=raycaster.cfg.max_distance + 1e2, + mesh_ids_wp=raycaster._mesh_ids_wp, + )[0].squeeze(0) + else: + hit_point = raycast_mesh( + ray_starts=self._height_grid_pos.unsqueeze(0), + ray_directions=direction.unsqueeze(0), + max_dist=raycaster.cfg.max_distance + 1e2, + mesh=raycaster.meshes[raycaster.cfg.mesh_prim_paths[0]], + )[0].squeeze(0) + + # get the height grid + self._height_grid = hit_point[:, 2].reshape( + int(np.abs(np.ceil((mesh_dimensions[0] - mesh_dimensions[2]) / self.cfg.grid_resolution))), + int(np.abs(np.ceil((mesh_dimensions[1] - mesh_dimensions[3]) / self.cfg.grid_resolution))), + ) + + def _construct_traversability_map(self): + # Define Sobel filters for x and y directions + sobel_x = ( + torch.tensor([[-1.0, 0.0, 1.0], [-2.0, 0.0, 2.0], [-1.0, 0.0, 1.0]], device=self.device) + .unsqueeze(0) + .unsqueeze(0) + ) + + sobel_y = ( + torch.tensor([[-1.0, -2.0, -1.0], [0.0, 0.0, 0.0], [1.0, 2.0, 1.0]], device=self.device) + .unsqueeze(0) + .unsqueeze(0) + ) + + # Apply the Sobel filters to the heigt_field_matrix while retraining the same shape + edges_x = torch.nn.functional.conv2d(self._height_grid.unsqueeze(0).float(), sobel_x, padding=1) + edges_y = torch.nn.functional.conv2d(self._height_grid.unsqueeze(0).float(), sobel_y, padding=1) + + # Compute the gradient magnitude (edge strength) + edges = torch.sqrt(edges_x**2 + edges_y**2) + + edges_mask = edges > self.cfg.edge_threshold + + # Dilate the mask to expand the objects + padding_size = int(self.cfg.robot_length / 2 / self.cfg.grid_resolution) + kernel = torch.ones((1, 1, 2 * padding_size + 1, 2 * padding_size + 1), device=self.device) + traversability_map = ( + torch.nn.functional.conv2d(edges_mask.float(), kernel, padding=padding_size).squeeze(1) < 0.5 + ) + traversability_map = traversability_map.reshape(-1, 1) + + # split the grid into subgrids for each terrain origin + self._terrain_origins = self._env.scene.terrain.terrain_origins.reshape(-1, 3) + point_distances = torch.cdist(self._height_grid_pos[..., :2], self._terrain_origins[..., :2], p=2) + subgrid_ids = torch.argmin(point_distances, dim=1) + + # split the traversability map into the subgrids based on the subgrid_ids + split_traversability_map = [ + self._height_grid_pos[subgrid_ids == i][traversability_map[subgrid_ids == i].squeeze()] + for i in range(self._terrain_origins.shape[0]) + ] + + # make every snipped the same length by repeating prev. indexes + split_lengths = [len(subgrid) for subgrid in split_traversability_map] + self.split_max_length = int(np.median(split_lengths)) + + self._split_traversability_map = torch.concat( + [ + subgrid[torch.randint(0, max(split_lengths[idx], 1), (self.split_max_length,))].unsqueeze(0) + for idx, subgrid in enumerate(split_traversability_map) + ], + dim=0, + ) + + # for cells where there is no accessible points, randomly assign points of another cell + zero_split_lengths = np.array(split_lengths) == 0 + if np.sum(zero_split_lengths) > 0: + while True: + rand_cell_idx = torch.randint(0, len(split_traversability_map), (np.sum(zero_split_lengths),)) + if np.all(np.array(split_lengths)[rand_cell_idx] > 0): + break + self._split_traversability_map[zero_split_lengths] = self._split_traversability_map[rand_cell_idx] + + + """ + Visualization + """ + + def _set_debug_vis_impl(self, debug_vis: bool): + """Set the debug visualization for the command. + + Args: + debug_vis (bool): Whether to enable debug visualization. + """ + # create markers if necessary for the first time + # for each marker type check that the correct command properties exist eg. need spawn position for spawn marker + if debug_vis: + if not hasattr(self, "box_goal_visualizer"): + marker_cfg = CUBOID_MARKER_CFG.copy() + marker_cfg.prim_path = "/Visuals/Command/position_goal" + marker_cfg.markers["cuboid"].size = (0.3, 0.3, 0.3) + marker_cfg.markers["cuboid"].visual_material.diffuse_color = (1.0, 0.15, 0.0) + marker_cfg.markers["cuboid"].visual_material.roughness = 0.7 + marker_cfg.markers["cuboid"].visual_material.metallic = 1.0 + self.box_goal_visualizer = VisualizationMarkers(marker_cfg) + self.box_goal_visualizer.set_visibility(True) + if self.cfg.vis_line and not hasattr(self, "line_to_goal_visualiser"): + marker_cfg = CYLINDER_MARKER_CFG.copy() + marker_cfg.prim_path = "/Visuals/Command/line_to_goal" + marker_cfg.markers["cylinder"].height = 1 + marker_cfg.markers["cylinder"].radius = 0.05 + self.line_to_goal_visualiser = VisualizationMarkers(marker_cfg) + self.line_to_goal_visualiser.set_visibility(True) + else: + if hasattr(self, "box_goal_visualizer"): + self.box_goal_visualizer.set_visibility(False) + if self.cfg.vis_line and hasattr(self, "line_to_goal_visualiser"): + self.line_to_goal_visualiser.set_visibility(False) + + def _debug_vis_callback(self, event, env_ids: Sequence[int] | None = None): + """Callback function for the debug visualization.""" + if env_ids is None: + env_ids = slice(None) + + # update goal marker if it exists + self.box_goal_visualizer.visualize(self.pos_command_w[env_ids]) + + if self.cfg.vis_line: + # update the line marker + # calculate the difference vector between the robot root position and the goal position + # TODO @tasdep this assumes that robot.data.body_pos_w exists + difference = self.pos_command_w - self.robot.data.body_pos_w[:, 0, :3] + translations = self.robot.data.body_pos_w[:, 0, :3] + # calculate the scale of the arrow (Mx3) + difference_norm = torch.norm(difference, dim=1) + # translate half of the length along difference axis + translations += difference / 2 + # scale along x axis + scales = torch.vstack( + [difference_norm, torch.ones_like(difference_norm), torch.ones_like(difference_norm)] + ).T + # convert the difference vector to a quaternion + difference = torch.nn.functional.normalize(difference, dim=1) + x_vec = torch.tensor([1, 0, 0]).float().to(self.pos_command_w.device) + angle = -torch.acos(difference @ x_vec) + axis = torch.linalg.cross(difference, x_vec.expand_as(difference)) + quat = quat_from_angle_axis(angle, axis) + + # apply transforms + self.line_to_goal_visualiser.visualize( + translations=translations[env_ids], scales=scales[env_ids], orientations=quat[env_ids] + ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/commands/goal_command_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/commands/goal_command_cfg.py new file mode 100644 index 00000000000..ad98ace71e7 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/commands/goal_command_cfg.py @@ -0,0 +1,49 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Code adapted from https://github.com/leggedrobotics/nav-suite + +# Copyright (c) 2025, The Nav-Suite Project Developers (https://github.com/leggedrobotics/nav-suite/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from dataclasses import MISSING + +from isaaclab.managers import CommandTermCfg +from isaaclab.utils import configclass + +from .goal_command import GoalCommandTerm + + +@configclass +class GoalCommandCfg(CommandTermCfg): + class_type: type = GoalCommandTerm + + vis_line: bool = True + """Whether to visualize the line from the robot to the goal.""" + + asset_name: str = "robot" + """Name of the asset in the environment for which the commands are generated.""" + + path_length_range: list[float] = [2.0, 10.0] + """Range of the sampled trajectories between start and goal.""" + + raycaster_sensor: str = MISSING + """Name of the raycaster sensor to use for terrain analysis.""" + + grid_resolution: float = 0.1 + """Resolution of the grid for the terrain analysis.""" + + robot_length: float = 1.0 + """Length of the robot for the terrain analysis.""" + + reset_pos_term_name: str = MISSING + """Name of the reset position term to use for the goal command.""" + + edge_threshold: float = 0.5 + """Threshold for the edge detection to define as untraversable.""" diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/events.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/events.py new file mode 100644 index 00000000000..1221f6e0e75 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/events.py @@ -0,0 +1,92 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Code adapted from https://github.com/leggedrobotics/nav-suite + +# Copyright (c) 2025, The Nav-Suite Project Developers (https://github.com/leggedrobotics/nav-suite/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import torch +from typing import TYPE_CHECKING + +from isaaclab.assets import Articulation +from isaaclab.managers import SceneEntityCfg +from isaaclab.utils.math import quat_from_euler_xyz, sample_uniform + +from .commands import GoalCommandTerm + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedRLEnv + + +def reset_robot_position( + env: ManagerBasedRLEnv, + env_ids: torch.Tensor, + yaw_range: tuple[float, float], + velocity_range: dict[str, tuple[float, float]] = {}, + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), + goal_command_generator_name: str = "goal_command", + spawn_in_env_frame: bool = False, + set_default_joint_state: bool = True, + add_default_base_pos: bool = False, +): + """Reset the asset root state to the spawn state defined by the command generator. + + Args: + env: The environment object. + env_ids: The environment ids to reset. + yaw_range: The additive heading range to apply to the spawn heading. + velocity_range: The velocity range to apply to the spawn velocity. + asset_cfg: The asset configuration to reset. Defaults to SceneEntityCfg("robot"). + spawn_in_env_frame: Whether spawn is based on the environment frame. Defaults to False. + set_default_joint_state: Whether to set the default joint state. Defaults to True. + add_default_base_pos: Whether to add the default base position to the spawn position. Defaults to False. + This is typically not necessary, as the robot height is added to the points in the + :class:`nav_collectors.terrain_analysis.TerrainAnalysis` Module + + .. note:: + This event assumes the existence of 'pos_spawn_w' and 'heading_spawn_w' in the command generator term, + i.e. an goal command generator of type :class:`nav_tasks.mdp.GoalCommand` is required. + See :class:`nav_tasks.mdp.TerrainAnalysisRootReset` for a more general approach. + + """ + # extract the used quantities (to enable type-hinting) + asset: Articulation = env.scene[asset_cfg.name] + goal_cmd_generator: GoalCommandTerm = env.command_manager.get_term(goal_command_generator_name) # type: ignore + + # positions - based on given start points (command generator) + positions = goal_cmd_generator.pos_spawn_w[env_ids] + if add_default_base_pos: + positions += asset.data.default_root_state[env_ids, :3] + if spawn_in_env_frame: + positions += env.scene.env_origins[env_ids] + + # yaw range + yaw_samples = sample_uniform(yaw_range[0], yaw_range[1], (len(env_ids), 1), device=asset.device) + yaw_samples += goal_cmd_generator.heading_spawn_w[env_ids].unsqueeze(1) + orientations = quat_from_euler_xyz( + torch.zeros_like(yaw_samples), torch.zeros_like(yaw_samples), yaw_samples + ).squeeze(1) + + # velocities - zero + range_list = [velocity_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]] + ranges = torch.tensor(range_list, device=asset.device) + rand_samples = sample_uniform(ranges[:, 0], ranges[:, 1], (len(env_ids), 6), device=asset.device) + velocities = asset.data.default_root_state[env_ids, 7:13] + rand_samples + + # set into the physics simulation + asset.write_root_pose_to_sim(torch.cat([positions, orientations], dim=-1), env_ids=env_ids) + asset.write_root_velocity_to_sim(velocities, env_ids=env_ids) + + if set_default_joint_state: + # obtain default joint positions + default_joint_pos = asset.data.default_joint_pos[env_ids].clone() + default_joint_vel = asset.data.default_joint_vel[env_ids].clone() + # set into the physics simulation + asset.write_joint_state_to_sim(default_joint_pos, default_joint_vel, env_ids=env_ids) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/observations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/observations.py new file mode 100644 index 00000000000..5fe765da548 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/observations.py @@ -0,0 +1,98 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Copyright (c) 2025, The Nav-Suite Project Developers (https://github.com/leggedrobotics/nav-suite/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import torch +from typing import TYPE_CHECKING + +from isaaclab.managers import SceneEntityCfg +from isaaclab.sensors import Camera, RayCasterCamera, TiledCamera + +from .actions import NavigationSE2Action + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedEnv, ManagerBasedRLEnv + + +""" +Sensors. +""" + + +def camera_image( + env: ManagerBasedEnv, + sensor_cfg: SceneEntityCfg, + data_type: str = "distance_to_image_plane", + flatten: bool = False, + nan_fill_value: float | None = None, +) -> torch.Tensor: + """Camera image Observations. + + The camera image observation from the given sensor w.r.t. the asset's root frame. + Also removes nan/inf values and sets them to the maximum distance of the sensor + + Args: + env: The environment object. + sensor_cfg: The name of the sensor. + data_type: The type of data to extract from the sensor. Default is "distance_to_image_plane". + flatten: If True, the image will be flattened to 1D. Default is False. + nan_fill_value: The value to fill nan/inf values with. If None, the maximum distance of the sensor will be used. + + Returns: + The image data.""" + # extract the used quantities (to enable type-hinting) + sensor: Camera | RayCasterCamera | TildCamera = env.scene.sensors[sensor_cfg.name] + + img = sensor.data.output[data_type].clone() + + if data_type == "distance_to_image_plane": + if nan_fill_value is None: + nan_fill_value = ( + sensor.cfg.max_distance if isinstance(sensor, RayCasterCamera) else sensor.cfg.spawn.clipping_range[1] + ) + img = torch.nan_to_num(img, nan=nan_fill_value, posinf=nan_fill_value, neginf=0.0) + + # if type torch.uint8, convert to float and scale between 0 and 1 + if img.dtype == torch.uint8: + img = img.to(torch.float32) / 255.0 + + if flatten: + return img.flatten(start_dim=1) + else: + # reorder the image to [BS, C, H, W] if it is not already in that shape + if img.shape[-1] == 1 or img.shape[-1] == 3: + img = img.permute(0, 3, 1, 2) + + return img + + +""" +Actions. +""" + + +def last_low_level_action( + env: ManagerBasedRLEnv, action_term: str, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot") +) -> torch.Tensor: + """The last low-level action.""" + action_term: NavigationSE2Action = env.action_manager._terms[action_term] + return action_term.low_level_actions[:, asset_cfg.joint_ids] + + +""" +Commands. +""" + + +def vel_commands(env: ManagerBasedRLEnv, action_term: str) -> torch.Tensor: + """The velocity command generated by the planner and given as input to the step function""" + action_term: NavigationSE2Action = env.action_manager._terms[action_term] + return action_term.processed_actions diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/pre_trained_policy_action.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/pre_trained_policy_action.py deleted file mode 100644 index 0437963698c..00000000000 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/pre_trained_policy_action.py +++ /dev/null @@ -1,188 +0,0 @@ -# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import torch -from dataclasses import MISSING -from typing import TYPE_CHECKING - -import isaaclab.utils.math as math_utils -from isaaclab.assets import Articulation -from isaaclab.managers import ActionTerm, ActionTermCfg, ObservationGroupCfg, ObservationManager -from isaaclab.markers import VisualizationMarkers -from isaaclab.markers.config import BLUE_ARROW_X_MARKER_CFG, GREEN_ARROW_X_MARKER_CFG -from isaaclab.utils import configclass -from isaaclab.utils.assets import check_file_path, read_file - -if TYPE_CHECKING: - from isaaclab.envs import ManagerBasedRLEnv - - -class PreTrainedPolicyAction(ActionTerm): - r"""Pre-trained policy action term. - - This action term infers a pre-trained policy and applies the corresponding low-level actions to the robot. - The raw actions correspond to the commands for the pre-trained policy. - - """ - - cfg: PreTrainedPolicyActionCfg - """The configuration of the action term.""" - - def __init__(self, cfg: PreTrainedPolicyActionCfg, env: ManagerBasedRLEnv) -> None: - # initialize the action term - super().__init__(cfg, env) - - self.robot: Articulation = env.scene[cfg.asset_name] - - # load policy - if not check_file_path(cfg.policy_path): - raise FileNotFoundError(f"Policy file '{cfg.policy_path}' does not exist.") - file_bytes = read_file(cfg.policy_path) - self.policy = torch.jit.load(file_bytes).to(env.device).eval() - - self._raw_actions = torch.zeros(self.num_envs, self.action_dim, device=self.device) - - # prepare low level actions - self._low_level_action_term: ActionTerm = cfg.low_level_actions.class_type(cfg.low_level_actions, env) - self.low_level_actions = torch.zeros(self.num_envs, self._low_level_action_term.action_dim, device=self.device) - - def last_action(): - # reset the low level actions if the episode was reset - if hasattr(env, "episode_length_buf"): - self.low_level_actions[env.episode_length_buf == 0, :] = 0 - return self.low_level_actions - - # remap some of the low level observations to internal observations - cfg.low_level_observations.actions.func = lambda dummy_env: last_action() - cfg.low_level_observations.actions.params = dict() - cfg.low_level_observations.velocity_commands.func = lambda dummy_env: self._raw_actions - cfg.low_level_observations.velocity_commands.params = dict() - - # add the low level observations to the observation manager - self._low_level_obs_manager = ObservationManager({"ll_policy": cfg.low_level_observations}, env) - - self._counter = 0 - - """ - Properties. - """ - - @property - def action_dim(self) -> int: - return 3 - - @property - def raw_actions(self) -> torch.Tensor: - return self._raw_actions - - @property - def processed_actions(self) -> torch.Tensor: - return self.raw_actions - - """ - Operations. - """ - - def process_actions(self, actions: torch.Tensor): - self._raw_actions[:] = actions - - def apply_actions(self): - if self._counter % self.cfg.low_level_decimation == 0: - low_level_obs = self._low_level_obs_manager.compute_group("ll_policy") - self.low_level_actions[:] = self.policy(low_level_obs) - self._low_level_action_term.process_actions(self.low_level_actions) - self._counter = 0 - self._low_level_action_term.apply_actions() - self._counter += 1 - - """ - Debug visualization. - """ - - def _set_debug_vis_impl(self, debug_vis: bool): - # set visibility of markers - # note: parent only deals with callbacks. not their visibility - if debug_vis: - # create markers if necessary for the first time - if not hasattr(self, "base_vel_goal_visualizer"): - # -- goal - marker_cfg = GREEN_ARROW_X_MARKER_CFG.copy() - marker_cfg.prim_path = "/Visuals/Actions/velocity_goal" - marker_cfg.markers["arrow"].scale = (0.5, 0.5, 0.5) - self.base_vel_goal_visualizer = VisualizationMarkers(marker_cfg) - # -- current - marker_cfg = BLUE_ARROW_X_MARKER_CFG.copy() - marker_cfg.prim_path = "/Visuals/Actions/velocity_current" - marker_cfg.markers["arrow"].scale = (0.5, 0.5, 0.5) - self.base_vel_visualizer = VisualizationMarkers(marker_cfg) - # set their visibility to true - self.base_vel_goal_visualizer.set_visibility(True) - self.base_vel_visualizer.set_visibility(True) - else: - if hasattr(self, "base_vel_goal_visualizer"): - self.base_vel_goal_visualizer.set_visibility(False) - self.base_vel_visualizer.set_visibility(False) - - def _debug_vis_callback(self, event): - # check if robot is initialized - # note: this is needed in-case the robot is de-initialized. we can't access the data - if not self.robot.is_initialized: - return - # get marker location - # -- base state - base_pos_w = self.robot.data.root_pos_w.clone() - base_pos_w[:, 2] += 0.5 - # -- resolve the scales and quaternions - vel_des_arrow_scale, vel_des_arrow_quat = self._resolve_xy_velocity_to_arrow(self.raw_actions[:, :2]) - vel_arrow_scale, vel_arrow_quat = self._resolve_xy_velocity_to_arrow(self.robot.data.root_lin_vel_b[:, :2]) - # display markers - self.base_vel_goal_visualizer.visualize(base_pos_w, vel_des_arrow_quat, vel_des_arrow_scale) - self.base_vel_visualizer.visualize(base_pos_w, vel_arrow_quat, vel_arrow_scale) - - """ - Internal helpers. - """ - - def _resolve_xy_velocity_to_arrow(self, xy_velocity: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - """Converts the XY base velocity command to arrow direction rotation.""" - # obtain default scale of the marker - default_scale = self.base_vel_goal_visualizer.cfg.markers["arrow"].scale - # arrow-scale - arrow_scale = torch.tensor(default_scale, device=self.device).repeat(xy_velocity.shape[0], 1) - arrow_scale[:, 0] *= torch.linalg.norm(xy_velocity, dim=1) * 3.0 - # arrow-direction - heading_angle = torch.atan2(xy_velocity[:, 1], xy_velocity[:, 0]) - zeros = torch.zeros_like(heading_angle) - arrow_quat = math_utils.quat_from_euler_xyz(zeros, zeros, heading_angle) - # convert everything back from base to world frame - base_quat_w = self.robot.data.root_quat_w - arrow_quat = math_utils.quat_mul(base_quat_w, arrow_quat) - - return arrow_scale, arrow_quat - - -@configclass -class PreTrainedPolicyActionCfg(ActionTermCfg): - """Configuration for pre-trained policy action term. - - See :class:`PreTrainedPolicyAction` for more details. - """ - - class_type: type[ActionTerm] = PreTrainedPolicyAction - """ Class of the action term.""" - asset_name: str = MISSING - """Name of the asset in the environment for which the commands are generated.""" - policy_path: str = MISSING - """Path to the low level policy (.pt files).""" - low_level_decimation: int = 4 - """Decimation factor for the low level action term.""" - low_level_actions: ActionTermCfg = MISSING - """Low level action configuration.""" - low_level_observations: ObservationGroupCfg = MISSING - """Low level observation configuration.""" - debug_vis: bool = True - """Whether to visualize debug information. Defaults to False.""" diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/rewards.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/rewards.py index 59c7ec5a936..62906dba8c1 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/rewards.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/rewards.py @@ -3,25 +3,60 @@ # # SPDX-License-Identifier: BSD-3-Clause +# Code adapted from https://github.com/leggedrobotics/nav-suite + +# Copyright (c) 2025, The Nav-Suite Project Developers (https://github.com/leggedrobotics/nav-suite/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + from __future__ import annotations import torch from typing import TYPE_CHECKING +from isaaclab.assets import Articulation +from isaaclab.managers import SceneEntityCfg + if TYPE_CHECKING: from isaaclab.envs import ManagerBasedRLEnv -def position_command_error_tanh(env: ManagerBasedRLEnv, std: float, command_name: str) -> torch.Tensor: - """Reward position tracking with tanh kernel.""" - command = env.command_manager.get_command(command_name) - des_pos_b = command[:, :3] - distance = torch.norm(des_pos_b, dim=1) - return 1 - torch.tanh(distance / std) +def backwards_movement(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: + """Reward the agent for moving backwards using L2-Kernel + + Args: + env: The learning environment. + asset_cfg: The name of the robot asset. + + Returns: + Dense reward [0, +1] based on the backward velocity. + """ + asset: Articulation = env.scene[asset_cfg.name] + # compute the reward + forward_velocity = asset.data.root_lin_vel_b[:, 0] + backward_movement_idx = torch.where( + forward_velocity < 0.0, torch.ones_like(forward_velocity), torch.zeros_like(forward_velocity) + ) + reward = torch.square(backward_movement_idx * forward_velocity) + reward = torch.clip(reward, min=0, max=1.0) + return reward + + +def lateral_movement(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: + """ + Reward the agent for moving lateral using L2-Kernel + Args: + env: The learning environment. + asset_cfg: The name of the robot asset. -def heading_command_error_abs(env: ManagerBasedRLEnv, command_name: str) -> torch.Tensor: - """Penalize tracking orientation error.""" - command = env.command_manager.get_command(command_name) - heading_b = command[:, 3] - return heading_b.abs() + Returns: + Dense reward [0, +1] based on the lateral velocity. + """ + asset: Articulation = env.scene[asset_cfg.name] + # compute the reward + lateral_velocity = asset.data.root_lin_vel_b[:, 1] + reward = torch.square(lateral_velocity) + reward = torch.clip(reward, min=0, max=1.0) + return reward diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/terminations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/terminations.py new file mode 100644 index 00000000000..1d03763dfe2 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/terminations.py @@ -0,0 +1,50 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Code adapted from https://github.com/leggedrobotics/nav-suite + +# Copyright (c) 2025, The Nav-Suite Project Developers (https://github.com/leggedrobotics/nav-suite/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import torch +from typing import TYPE_CHECKING + +from isaaclab.assets import Articulation +from isaaclab.managers import SceneEntityCfg + +from .commands import GoalCommandTerm + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedRLEnv + + +def at_goal( + env: ManagerBasedRLEnv, + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), + distance_threshold: float = 0.5, + command_generator_term_name: str = "goal_command", +) -> torch.Tensor: + """Terminate the episode when the goal is reached. + + Args: + env: The learning environment. + asset_cfg: The name of the robot asset. + distance_threshold: The distance threshold to the goal. + speed_threshold: The speed threshold at the goal. + + Returns: + Boolean tensor indicating whether the goal is reached. + """ + # extract the used quantities (to enable type-hinting) + asset: Articulation = env.scene[asset_cfg.name] + goal_cmd_generator: GoalCommandTerm = env.command_manager.get_term(command_generator_term_name) + + # Check conditions for termination + distance_goal = torch.norm(asset.data.root_pos_w[:, :2] - goal_cmd_generator.pos_command_w[:, :2], dim=1, p=2) + return distance_goal < distance_threshold diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/navigation_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/navigation_env_cfg.py new file mode 100644 index 00000000000..9dfa94129d1 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/navigation_env_cfg.py @@ -0,0 +1,408 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Code adapted from https://github.com/leggedrobotics/nav-suite + +# Copyright (c) 2025, The Nav-Suite Project Developers (https://github.com/leggedrobotics/nav-suite/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import os + +# Set the PYTORCH_CUDA_ALLOC_CONF environment variable +os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:256" + +import isaaclab.sim as sim_utils +from isaaclab.envs import ManagerBasedRLEnvCfg +from isaaclab.managers import EventTermCfg as EventTerm +from isaaclab.managers import ObservationGroupCfg as ObsGroup +from isaaclab.managers import ObservationTermCfg as ObsTerm +from isaaclab.managers import RewardTermCfg as RewTerm +from isaaclab.managers import SceneEntityCfg +from isaaclab.managers import TerminationTermCfg as DoneTerm +from isaaclab.sensors import RayCasterCameraCfg, TiledCameraCfg, patterns +from isaaclab.terrains import TerrainGeneratorCfg +from isaaclab.terrains.height_field import HfRandomUniformTerrainCfg +from isaaclab.utils import configclass +from isaaclab.utils.assets import ISAACLAB_NUCLEUS_DIR + +import isaaclab_tasks.manager_based.locomotion.velocity.velocity_env_cfg as LOW_LEVEL_CFGS +import isaaclab_tasks.manager_based.navigation.mdp as mdp + +from .terrains import MeshPillarTerrainCfg + +## +# Scene definition +## + + +NAV_TERRAIN = TerrainGeneratorCfg( + size=(8.0, 8.0), + border_width=1.0, + border_height=-3.0, + num_rows=10, + num_cols=20, + horizontal_scale=0.1, + vertical_scale=0.005, + slope_threshold=0.75, + use_cache=False, + sub_terrains={ + "pillar": MeshPillarTerrainCfg( + proportion=0.8, + box_objects=MeshPillarTerrainCfg.BoxCfg( + width=(0.2, 0.4), + length=(0.2, 0.4), + height=(0.4, 1.5), + num_objects=(5, 15), + max_yx_angle=(0.0, 20.0), + ), + cylinder_cfg=MeshPillarTerrainCfg.CylinderCfg( + radius=(0.1, 0.3), + height=(0.4, 1.5), + num_objects=(5, 15), + max_yx_angle=(0.0, 20.0), + ), + ), + "pillar_road": MeshPillarTerrainCfg( + proportion=0.2, + box_objects=MeshPillarTerrainCfg.BoxCfg( + width=(0.2, 0.4), + length=(0.2, 0.4), + height=(0.2, 0.4), + num_objects=(10, 20), + ), + cylinder_cfg=MeshPillarTerrainCfg.CylinderCfg( + radius=(0.1, 0.3), + height=(0.4, 1.5), + num_objects=(5, 15), + max_yx_angle=(0.0, 20.0), + ), + rough_terrain=HfRandomUniformTerrainCfg( + proportion=0.2, noise_range=(0.02, 0.10), noise_step=0.02, border_width=0.25 + ), + ), + }, +) + + +@configclass +class RayCasterNavSceneCfg(LOW_LEVEL_CFGS.MySceneCfg): + """Configuration for a scene for training a perceptive navigation policy on an AnymalD Robot.""" + + # SENSORS: Navigation Policy + front_camera = RayCasterCameraCfg( + prim_path="{ENV_REGEX_NS}/Robot/base", + mesh_prim_paths=["/World/ground"], + pattern_cfg=patterns.PinholeCameraPatternCfg( + height=36, + width=64, + ), + update_period=0, + debug_vis=False, + offset=RayCasterCameraCfg.OffsetCfg( + pos=(0.4761, 0.0035, 0.1055), + rot=(0.9914449, 0.0, 0.1305262, 0.0), + convention="world", # 15 degrees downward tilted + ), + max_distance=20, + data_types=["distance_to_image_plane"], + ) + + def __post_init__(self): + """Post initialization.""" + # swap to navigation terrain + self.terrain.terrain_generator = NAV_TERRAIN + + +@configclass +class TiledNavSceneCfg(LOW_LEVEL_CFGS.MySceneCfg): + """Configuration for a scene for training a perceptive navigation policy on an AnymalD Robot.""" + + # SENSORS: Navigation Policy + front_camera = TiledCameraCfg( + prim_path="{ENV_REGEX_NS}/Robot/base/camera", + spawn=sim_utils.PinholeCameraCfg( + clipping_range=(0.01, 20.0), + ), + height=36, + width=64, + update_period=0, + offset=TiledCameraCfg.OffsetCfg( + pos=(0.4761, 0.0035, 0.1055), + rot=(0.9914449, 0.0, 0.1305262, 0.0), + convention="world", # 15 degrees downward tilted + ), + data_types=["distance_to_image_plane"], + ) + + def __post_init__(self): + """Post initialization.""" + # swap to navigation terrain + self.terrain.terrain_generator = NAV_TERRAIN + + +## +# MDP settings +## + + +@configclass +class ActionsCfg: + """Action specifications for the MDP.""" + + velocity_command = mdp.NavigationSE2ActionCfg( + asset_name="robot", + low_level_action=LOW_LEVEL_CFGS.ActionsCfg().joint_pos, + low_level_policy_file=ISAACLAB_NUCLEUS_DIR + "/Policies/ANYmal-C/HeightScan/policy.pt", + clip_mode="minmax", + clip=[ + LOW_LEVEL_CFGS.CommandsCfg().base_velocity.ranges.lin_vel_x, + LOW_LEVEL_CFGS.CommandsCfg().base_velocity.ranges.lin_vel_y, + LOW_LEVEL_CFGS.CommandsCfg().base_velocity.ranges.ang_vel_z, + ], + ) + + +@configclass +class NavObservationsCfg: + """Observation specifications for the MDP.""" + + @configclass + class NavProprioceptiveCfg(ObsGroup): + """Proprioceptive observations for navigation policy group.""" + + base_lin_vel = ObsTerm(func=mdp.base_lin_vel) + base_ang_vel = ObsTerm(func=mdp.base_ang_vel) + joint_pos = ObsTerm(func=mdp.joint_pos_rel) + joint_vel = ObsTerm(func=mdp.joint_vel_rel) + + goal_commands = ObsTerm(func=mdp.generated_commands, params={"command_name": "goal_command"}) + + def __post_init__(self): + self.enable_corruption = True + self.concatenate_terms = True + + @configclass + class NavExteroceptiveCfg(ObsGroup): + """Exteroceptive observations for navigation policy group.""" + + forwards_depth_image = ObsTerm( + func=mdp.camera_image, + params={"sensor_cfg": SceneEntityCfg("front_camera")}, + ) + + def __post_init__(self): + self.enable_corruption = True + self.concatenate_terms = True + + # Observation Groups + low_level_policy: LOW_LEVEL_CFGS.ObservationsCfg.PolicyCfg = LOW_LEVEL_CFGS.ObservationsCfg.PolicyCfg() + proprioceptive: NavProprioceptiveCfg = NavProprioceptiveCfg() + exteroceptive: NavExteroceptiveCfg = NavExteroceptiveCfg() + + def __post_init__(self): + # adjust because the velocity commands are now given by the navigation policy + self.low_level_policy.velocity_commands = ObsTerm( + func=mdp.vel_commands, params={"action_term": "velocity_command"} + ) + self.low_level_policy.actions = ObsTerm( + func=mdp.last_low_level_action, params={"action_term": "velocity_command"} + ) + + +@configclass +class EventCfg: + """Configuration for randomization.""" + + reset_base = EventTerm( + func=mdp.reset_robot_position, + mode="reset", + params={ + "asset_cfg": SceneEntityCfg("robot"), + "yaw_range": (-3.14, 3.14), + "velocity_range": { + "x": (-0.5, 0.5), + "y": (-0.5, 0.5), + "z": (0, 0), + "roll": (0, 0), + "pitch": (0, 0), + "yaw": (-0.5, 0.5), + }, + "goal_command_generator_name": "goal_command", + }, + ) + + reset_robot_joints = EventTerm( + func=mdp.reset_joints_by_scale, + mode="reset", + params={ + "asset_cfg": SceneEntityCfg("robot"), + "position_range": (0.0, 0.0), + "velocity_range": (0.0, 0.0), + }, + ) + + +@configclass +class RewardsCfg: + """Reward terms for the MDP. + + .. note:: + All rewards get multiplied with weight*dt - consider this when setting weights. + Rewards are normalized over max episode length in wandb logging. + + .. note:: + In wandb: + - Episode Rewards are in seconds + - Train Mean Reward is based on episode length (Rewards * Episode Length) + """ + + # -- rewards + goal_reached_rew = RewTerm( + func=mdp.is_terminated_term, + params={"term_keys": "goal_reached"}, + weight=200.0, + ) + + # -- penalties + lateral_movement = RewTerm( + func=mdp.lateral_movement, + weight=-0.01, + ) + backward_movement = RewTerm( + func=mdp.backwards_movement, + weight=-0.01, + ) + episode_termination = RewTerm( + func=mdp.is_terminated_term, # type: ignore + params={"term_keys": ["base_contact"]}, + weight=-200.0, + ) + action_rate_l2 = RewTerm( + func=mdp.action_rate_l2, + weight=-0.1, + ) + + +@configclass +class TerminationsCfg: + """Termination terms for the MDP.""" + + time_out = DoneTerm(func=mdp.time_out, time_out=True) + + goal_reached = DoneTerm( + func=mdp.at_goal, # type: ignore + params={ + "distance_threshold": 0.5, + "command_generator_term_name": "goal_command", + }, + time_out=False, + ) + + base_contact = DoneTerm( + func=mdp.illegal_contact, + params={ + "sensor_cfg": SceneEntityCfg("contact_forces", body_names=["base"]), + "threshold": 0.0, + }, + time_out=False, + ) + + +@configclass +class CommandsCfg: + """Command specifications for the MDP.""" + + goal_command = mdp.GoalCommandCfg( + asset_name="robot", + grid_resolution=0.1, + robot_length=1.0, + raycaster_sensor="height_scanner", + resampling_time_range=(1.0e9, 1.0e9), # No resampling + debug_vis=True, + reset_pos_term_name="reset_base", + ) + + +## +# Environment configuration +## + + +@configclass +class NavEnvCfg(ManagerBasedRLEnvCfg): + """Configuration for the navigation environment.""" + + # Basic settings + observations: NavObservationsCfg = NavObservationsCfg() + actions: ActionsCfg = ActionsCfg() + commands: CommandsCfg = CommandsCfg() + + # MDP settings + rewards: RewardsCfg = RewardsCfg() + terminations: TerminationsCfg = TerminationsCfg() + events: EventCfg = EventCfg() + + def __post_init__(self): + """Post initialization.""" + + # Simulation settings + self.sim.dt = 0.005 # In seconds + self.sim.disable_contact_processing = True + self.sim.physics_material = self.scene.terrain.physics_material + + # General settings + self.episode_length_s = 20 + + # This sets how many times the high-level actions (navigation policy) + # are applied to the sim before being recalculated. + self.decimation = int(1 / self.sim.dt / 10) # 10Hz planner frequency + + # Similar to above, the low-level actions (locomotion controller) are calculated every: + # self.sim.dt * self.low_level_decimation, so 0.005 * 4 = 0.02 seconds, or 50Hz. + self.low_level_decimation = 4 + + +@configclass +class RayCasterNavEnvCfg(NavEnvCfg): + """Configuration for the navigation environment with ray caster camera.""" + + scene: RayCasterNavSceneCfg = RayCasterNavSceneCfg(num_envs=100, env_spacing=8) + + def __post_init__(self): + """Post initialization.""" + super().__post_init__() + + # update sensor update periods + # We tick contact sensors based on the smallest update period (physics update period) + if self.scene.contact_forces is not None: + self.scene.contact_forces.update_period = self.sim.dt + + # We tick the cameras based on the navigation policy update period. + if self.scene.front_camera is not None: + self.scene.front_camera.update_period = self.decimation * self.sim.dt + + +@configclass +class TiledNavEnvCfg(NavEnvCfg): + """Configuration for the navigation environment with tiled camera.""" + + scene: TiledNavSceneCfg = TiledNavSceneCfg(num_envs=100, env_spacing=8) + + def __post_init__(self): + """Post initialization.""" + super().__post_init__() + + # update sensor update periods + # We tick contact sensors based on the smallest update period (physics update period) + if self.scene.contact_forces is not None: + self.scene.contact_forces.update_period = self.sim.dt + + # We tick the cameras based on the navigation policy update period. + if self.scene.front_camera is not None: + self.scene.front_camera.update_period = self.decimation * self.sim.dt diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/terrains/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/terrains/__init__.py new file mode 100644 index 00000000000..69d16b05dd5 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/terrains/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from .pillar_terrain_cfg import MeshPillarTerrainCfg diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/terrains/pillar_terrain.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/terrains/pillar_terrain.py new file mode 100644 index 00000000000..2354834a2c1 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/terrains/pillar_terrain.py @@ -0,0 +1,150 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Code adapted from https://github.com/leggedrobotics/nav-suite + +# Copyright (c) 2025, The Nav-Suite Project Developers (https://github.com/leggedrobotics/nav-suite/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Functions to generate different terrains using the ``trimesh`` library.""" + +from __future__ import annotations + +import numpy as np +import trimesh +from typing import TYPE_CHECKING + +from isaaclab.terrains.height_field.hf_terrains import random_uniform_terrain +from isaaclab.terrains.trimesh.utils import * # noqa: F401, F403 +from isaaclab.terrains.trimesh.utils import make_plane + +if TYPE_CHECKING: + from . import pillar_terrain_cfg + + +def pillar_terrain( + difficulty: float, cfg: pillar_terrain_cfg.MeshPillarTerrainCfg +) -> tuple[list[trimesh.Trimesh], np.ndarray]: + """Generate a terrain with a set of repeated boxes and cylinders. + + .. image:: ../../_static/terrains/mesh_pillar_terrain.jpg + :width: 45% + :align: center + + The terrain has a ground with a platform in the middle. The objects are randomly placed on the + terrain s.t. they do not overlap with the platform. + + Args: + difficulty: The difficulty of the terrain. This is a value between 0 and 1. + cfg: The configuration for the terrain. + + Returns: + A tuple containing the tri-mesh of the terrain and the origin of the terrain (in m). + + Raises: + ValueError: If the object type is not supported. It must be either a string or a callable. + """ + from .pillar_terrain_cfg import MeshPillarTerrainCfg + + # initialize list of meshes + meshes_list = list() + # initialize list of object meshes + object_center_list = list() + # constants for the terrain + platform_clearance = 0.1 + # compute quantities + origin = np.asarray((0.5 * cfg.size[0], 0.5 * cfg.size[1], 0)) + platform_corners = np.asarray([ + [origin[0] - cfg.platform_width / 2, origin[1] - cfg.platform_width / 2], + [origin[0] + cfg.platform_width / 2, origin[1] + cfg.platform_width / 2], + ]) + platform_corners[0, :] *= 1 - platform_clearance + platform_corners[1, :] *= 1 + platform_clearance + # generate a ground plane for the terrain + ground_plane = make_plane(cfg.size, height=0.0, center_zero=False) + meshes_list.append(ground_plane) + # create rough terrain + if cfg.rough_terrain: + cfg.rough_terrain.size = cfg.size + rough_mesh, _ = random_uniform_terrain(difficulty, cfg.rough_terrain) + meshes_list += rough_mesh + + for object_cfg in [cfg.box_objects, cfg.cylinder_cfg]: + # if object type is a string, get the function: make_{object_type} + if isinstance(object_cfg.object_type, str): + object_func = globals().get(f"make_{object_cfg.object_type}") + else: + object_func = object_cfg.object_type + if not callable(object_func): + raise ValueError(f"The attribute 'object_type' must be a string or a callable. Received: {object_func}") + + # Resolve the terrain configuration + # -- common parameters + num_objects = object_cfg.num_objects[0] + int( + difficulty * (object_cfg.num_objects[1] - object_cfg.num_objects[0]) + ) + height = object_cfg.height[0] + difficulty * (object_cfg.height[1] - object_cfg.height[0]) + # -- object specific parameters + # note: SIM114 requires duplicated logical blocks under a single body. + if isinstance(object_cfg, MeshPillarTerrainCfg.BoxCfg): + object_kwargs = { + "length": object_cfg.length[0] + difficulty * (object_cfg.length[1] - object_cfg.length[0]), + "width": object_cfg.width[0] + difficulty * (object_cfg.width[1] - object_cfg.width[0]), + "max_yx_angle": object_cfg.max_yx_angle[0] + difficulty * ( + object_cfg.max_yx_angle[1] - object_cfg.max_yx_angle[0] + ), + "degrees": object_cfg.degrees, + } + elif isinstance(object_cfg, MeshPillarTerrainCfg.CylinderCfg): # noqa: SIM114 + object_kwargs = { + "radius": object_cfg.radius[0] + difficulty * (object_cfg.radius[1] - object_cfg.radius[0]), + "max_yx_angle": object_cfg.max_yx_angle[0] + difficulty * ( + object_cfg.max_yx_angle[1] - object_cfg.max_yx_angle[0] + ), + "degrees": object_cfg.degrees, + } + else: + raise ValueError(f"Unknown terrain configuration: {cfg}") + + # sample center for objects + while True: + object_centers = np.zeros((num_objects, 3)) + object_centers[:, 0] = np.random.uniform(0, cfg.size[0], num_objects) + object_centers[:, 1] = np.random.uniform(0, cfg.size[1], num_objects) + # filter out the centers that are on the platform + is_within_platform_x = np.logical_and( + object_centers[:, 0] >= platform_corners[0, 0], object_centers[:, 0] <= platform_corners[1, 0] + ) + is_within_platform_y = np.logical_and( + object_centers[:, 1] >= platform_corners[0, 1], object_centers[:, 1] <= platform_corners[1, 1] + ) + masks = np.logical_and(is_within_platform_x, is_within_platform_y) + # if there are no objects on the platform, break + if not np.any(masks): + break + + # generate obstacles (but keep platform clean) + for index in range(len(object_centers)): + # randomize the height of the object + ob_height = height + np.random.uniform(-cfg.max_height_noise, cfg.max_height_noise) + object_centers[index, 2] = ob_height / 2 + if ob_height > 0.0: + object_mesh = object_func(center=object_centers[index], height=ob_height, **object_kwargs) + meshes_list.append(object_mesh) + + object_center_list.append(object_centers) + + # generate a platform in the middle + dim = (cfg.platform_width, cfg.platform_width, 0) + pos = (0.5 * cfg.size[0], 0.5 * cfg.size[1], 0) + platform = trimesh.creation.box(dim, trimesh.transformations.translation_matrix(pos)) + meshes_list.append(platform) + + # elevate origin + # origin[2] = mean_height + + return meshes_list, origin diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/terrains/pillar_terrain_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/terrains/pillar_terrain_cfg.py new file mode 100644 index 00000000000..bfc42e1fa44 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/terrains/pillar_terrain_cfg.py @@ -0,0 +1,96 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Code adapted from https://github.com/leggedrobotics/nav-suite + +# Copyright (c) 2025, The Nav-Suite Project Developers (https://github.com/leggedrobotics/nav-suite/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from dataclasses import MISSING + +from isaaclab.terrains import SubTerrainBaseCfg +from isaaclab.terrains.height_field import HfRandomUniformTerrainCfg +from isaaclab.utils import configclass + +from .pillar_terrain import pillar_terrain + + +@configclass +class MeshPillarTerrainCfg(SubTerrainBaseCfg): + @configclass + class CylinderCfg: + """Configuration for repeated cylinder.""" + + radius: tuple[float, float] = MISSING + """The radius of the pyramids (in m). First value start of curriculum, second value end.""" + + max_yx_angle: tuple[float, float] = 0.0 + """The maximum angle along the y and x axis. Defaults to 0.0. First value start of curriculum, second value end.""" + + degrees: bool = True + """Whether the angle is in degrees. Defaults to True.""" + + num_objects: tuple[int, int] = MISSING + """The number of objects to add to the terrain. First value start of curriculum, second value end.""" + + height: tuple[float, float] = MISSING + """The height (along z) of the object (in m). First value start of curriculum, second value end.""" + + object_type: str = "cylinder" + """The type of object to generate. + + The type can be a string or a callable. If it is a string, the function will look for a function called + ``make_{object_type}`` in the current module scope. If it is a callable, the function will + use the callable to generate the object. + """ + + @configclass + class BoxCfg: + """Configuration for repeated boxes.""" + + width: tuple[float, float] = MISSING + """The width (along x) and length (along y) of the box (in m). First value start of curriculum, second value end.""" + + length: tuple[float, float] = MISSING + """The length (along y) of the box (in m). First value start of curriculum, second value end.""" + + max_yx_angle: tuple[float, float] = (0.0, 0.0) + """The maximum angle along the y and x axis. Defaults to (0.0, 0.0). First value start of curriculum, second value end.""" + + degrees: bool = True + """Whether the angle is in degrees. Defaults to True.""" + + num_objects: tuple[int, int] = MISSING + """The number of objects to add to the terrain. First value start of curriculum, second value end.""" + + height: tuple[float, float] = MISSING + """The height (along z) of the object (in m). First value start of curriculum, second value end.""" + + object_type: str = "box" + """The type of object to generate. + + The type can be a string or a callable. If it is a string, the function will look for a function called + ``make_{object_type}`` in the current module scope. If it is a callable, the function will + use the callable to generate the object. + """ + + function = pillar_terrain + + rough_terrain: HfRandomUniformTerrainCfg = None + """The configuration for the rough terrain. If None, the terrain will be flat. Defaults to None.""" + + box_objects: BoxCfg = MISSING + """add boxes to the terrain""" + + cylinder_cfg: CylinderCfg = MISSING + """add cylinders to the terrain""" + + max_height_noise: float = 0.0 + """The maximum amount of noise to add to the height of the objects (in m). Defaults to 0.0.""" + + platform_width: float = 1.0 + """The width of the cylindrical platform at the center of the terrain. Defaults to 1.0."""