diff --git a/docs/source/overview/core-concepts/sensors/camera.rst b/docs/source/overview/core-concepts/sensors/camera.rst index 19e1c6ee9742..8b9bc3358079 100644 --- a/docs/source/overview/core-concepts/sensors/camera.rst +++ b/docs/source/overview/core-concepts/sensors/camera.rst @@ -239,18 +239,21 @@ the project page: https://research.nvidia.com/labs/sil/projects/ppisp . PPISP is typically authored alongside a `ParticleField3DGaussianSplat `__ -USD asset: it carries a `RenderProduct -`__ -whose target camera and PPISP `UsdShade.Shader -`__ -(a shader prim named ``PPISP`` whose inputs follow the PPISP naming -convention) were calibrated against the real capture rig that produced -the splats. Configuring the camera with the matching PPISP coefficients -makes the rendered tile match the calibration target. - -The pipeline applies, in order: responsivity → exposure → vignetting → -color homography → camera response function → uint8 clamp. It runs as a -single Warp kernel. +USD asset: its camera prim authors ``ppisp:*`` attributes calibrated against +the real capture rig that produced the splats. Configuring the camera with +the matching PPISP coefficients makes the rendered tile match the calibration +target. + +With static coefficients, the image pass applies, in order: responsivity → +exposure → vignetting → color homography → camera response function → uint8 +clamp. It runs as a single Warp image kernel. + +When the camera also authors ``ppisp:controllerWeights``, Isaac Lab runs the +exported PPISP controller before the image pass. The controller reads the +current HDR image, prior exposure, and controller responsivity, then predicts +``exposureOffset`` and the four color-latent pairs for each camera view. The +image pass then uses those predicted exposure and color-latent values while +static PPISP inputs still provide responsivity, vignetting, and CRF. Configuration ^^^^^^^^^^^^^ @@ -259,15 +262,18 @@ Configuration * ``None`` (default) — ISP disabled. * :class:`~isaaclab_ppisp.PpispCfg` — explicit PPISP coefficients - (:attr:`~isaaclab_ppisp.PpispCfg.inputs`), or - :attr:`~isaaclab_ppisp.PpispCfg.shader_prim_path` to import them from a - PPISP ``UsdShade.Shader`` already on the stage. -* :class:`~sensors.CameraISPMode` — auto-discover an ISP shader on the + (:attr:`~isaaclab_ppisp.PpispCfg.inputs`) and, optionally, controller + weights. Use :attr:`~isaaclab_ppisp.PpispCfg.camera_prim_path` to import + static coefficients and camera-authored controller weights from a USD camera + already on the stage. +* :class:`~sensors.CameraISPMode` — auto-discover ISP camera attributes on the stage (see below). -The cfg applies once per Camera sensor batch. The PPISP Warp kernel takes -scalar coefficients, so every cloned view in a tiled batch shares the same -ISP configuration — there is no per-view ISP today. +The cfg applies once per Camera sensor batch. Static PPISP inputs are scalar +coefficients shared by every cloned view in a tiled batch. When controller +weights are present, the controller predicts per-view exposure and color-latent +values from each HDR image, while responsivity, vignetting, and CRF remain +shared by the batch. .. code-block:: python @@ -280,8 +286,8 @@ ISP configuration — there is no per-view ISP today. # explicit coefficients cfg = CameraCfg(..., isp_cfg=PpispCfg(inputs={"exposureOffset": 1.5})) - # import coefficients from a USD shader path - cfg = CameraCfg(..., isp_cfg=PpispCfg(shader_prim_path="/World/Render/PPISP")) + # import coefficients from a USD camera path + cfg = CameraCfg(..., isp_cfg=PpispCfg(camera_prim_path="/World/Camera_ppisp")) # auto-discover from the stage cfg = CameraCfg(..., isp_cfg=CameraISPMode.AUTO_ANY) @@ -293,22 +299,19 @@ Auto-discovery is opt-in via :class:`~sensors.CameraISPMode`. Discovery runs once at camera construction using the first matched camera prim in the Camera sensor batch: -1. Walk the stage for a USD ``RenderProduct`` whose ``camera`` relationship - targets the first matched camera prim **and** that has a child - ``UsdShade.Shader`` prim named ``PPISP``. If found, import its inputs as a - :class:`~isaaclab_ppisp.PpispCfg`. -2. ``AUTO_ANY`` only: if step 1 finds nothing, fall back to the first - ``UsdShade.Shader`` prim named ``PPISP`` anywhere on the stage. +1. Check the first matched camera prim for ``ppisp:*`` attributes. +2. ``AUTO_ANY`` only, or when no camera path is available: fall back to the + first camera anywhere on the stage with ``ppisp:*`` attributes. 3. Otherwise the ISP stays disabled for the whole Camera sensor batch. In practice this means: if the stage carries a ``ParticleField3DGaussianSplat`` -together with a ``RenderProduct`` that binds a ``PPISP`` shader child to the -batch's first matched camera prim, the Camera sensor picks up the matching ISP -automatically and no Python-side coefficient authoring is required. +together with a camera that authors ``ppisp:*`` attributes, the Camera sensor +picks up the matching ISP automatically, including controller weights when +authored, and no Python-side coefficient authoring is required. -``AUTO_CAMERA`` runs step 1 only — useful when the stage carries multiple -PPISP shaders and you want the Camera sensor batch to use exactly the one bound -to its first matched camera prim. +``AUTO_CAMERA`` runs only the camera-local discovery steps — useful when the +stage carries multiple PPISP cameras and you want the Camera sensor batch to use +the attributes authored on its first matched camera prim. Renderer support ^^^^^^^^^^^^^^^^ @@ -320,7 +323,8 @@ internally: the Isaac RTX renderer sources HDR from the Replicator Newton Warp renderer from its native scene-linear color buffer. Each backend allocates its own HDR scratch buffer when the user did not request ``"rgb_hdr"`` in :attr:`~sensors.CameraCfg.data_types`, and dispatches the -PPISP kernel into ``rgb`` / ``rgba`` after every render tick. +PPISP pipeline into ``rgb`` / ``rgba`` after every render tick. For controller +configs, this is the controller pass followed by the PPISP image pass. Usage example ^^^^^^^^^^^^^ @@ -337,7 +341,9 @@ absolute-difference images. Use ``--renderer isaac_rtx`` to run the same workflow with Isaac RTX. Pass ``--input_scene`` for a custom scene and ``--camera_prim_path`` if the stage -contains multiple PPISP-bound cameras. Images are written to +contains multiple cameras with PPISP attributes. If a config or command selects +a visualizer, force-disable all visualizers with ``--visualizer none`` or +``--viz none``. Images are written to ``scripts/demos/sensors/output/ppisp_camera`` unless ``--output_dir`` is set. Known limitations @@ -345,10 +351,11 @@ Known limitations * The ISP writes back into the ``rgb`` / ``rgba`` buffers. If neither is requested, configuring ``isp_cfg`` raises at camera init. -* PPISP inputs are static for the lifetime of the camera. Animated USD - shader inputs are collapsed to their first authored time sample. -* Coefficients are global per camera — no per-pixel or per-region - authoring beyond the radial vignetting term. +* Static PPISP inputs and controller weights are fixed for the lifetime of the + camera. Animated USD camera attributes are collapsed to their first authored + time sample. +* Static coefficients are global per Camera sensor batch — no per-pixel or + per-region authoring beyond the radial vignetting term. * PPISP is the only ISP implementation today. Other ISP families would need a new config type and discoverer entry. * On the Isaac RTX and OVRTX backends, enabling ``isp_cfg`` forces RTX-side @@ -356,8 +363,8 @@ Known limitations and authors a neutral ``OmniRtxCameraExposureAPI_1`` schema on each camera prim so the post-render ISP is the only path that processes color. Mixing this with RTX-side exposure authoring is not supported. -* Auto-discovery resolves at camera construction; later authoring of a - ``RenderProduct`` or shader on the stage is not picked up. +* Auto-discovery resolves at camera construction; later authoring of + ``ppisp:*`` camera attributes on the stage is not picked up. Depth and Distances ~~~~~~~~~~~~~~~~~~~ diff --git a/scripts/demos/sensors/ppisp_camera.py b/scripts/demos/sensors/ppisp_camera.py index 04a21dc6c1fe..2d609b354b28 100644 --- a/scripts/demos/sensors/ppisp_camera.py +++ b/scripts/demos/sensors/ppisp_camera.py @@ -42,7 +42,7 @@ "--camera_prim_path", type=str, default=None, - help="Optional camera prim path override. Omit to auto-select the first PPISP-bound camera.", + help="Optional camera prim path override. Omit to auto-select the first camera with PPISP attributes.", ) parser.add_argument( "--camera_time_code", @@ -126,9 +126,14 @@ import matplotlib.pyplot as plt import numpy as np import torch -from isaaclab_ppisp.cfg import PpispCfg, ppisp_cfg_from_usd_shader +from isaaclab_ppisp._demo_utils import ( + find_ppisp_camera_bindings, + format_available_ppisp_cameras, + order_ppisp_bindings_by_camera, +) +from isaaclab_ppisp.cfg import PpispCfg, ppisp_cfg_from_usd_camera -from pxr import Usd, UsdGeom, UsdShade +from pxr import Usd, UsdGeom import isaaclab.sim as sim_utils from isaaclab.assets import AssetBaseCfg, RigidObjectCfg @@ -191,61 +196,11 @@ def make_sim_cfg() -> sim_utils.SimulationCfg: ) -def find_ppisp_camera_bindings(stage: Usd.Stage) -> list[tuple[str, Usd.Prim, Usd.Prim]]: - """Return camera, RenderProduct, and PPISP shader bindings from a source stage.""" - bindings = [] - for prim in stage.Traverse(): - if prim.GetTypeName() != "RenderProduct": - continue - shader_prim = stage.GetPrimAtPath(prim.GetPath().AppendChild("PPISP")) - if not shader_prim or not shader_prim.IsValid(): - continue - camera_rel = prim.GetRelationship("camera") - if not camera_rel: - continue - for target in camera_rel.GetTargets(): - camera_prim = stage.GetPrimAtPath(target) - if camera_prim and camera_prim.IsValid(): - bindings.append((str(target), prim, shader_prim)) - return bindings - - -def order_ppisp_bindings_by_camera( - stage: Usd.Stage, ppisp_bindings: list[tuple[str, Usd.Prim, Usd.Prim]] -) -> list[tuple[str, Usd.Prim, Usd.Prim]]: - """Return PPISP bindings ordered by source Camera prim traversal.""" - binding_by_camera_path = {} - for binding in ppisp_bindings: - binding_by_camera_path.setdefault(binding[0], binding) - - ordered_bindings = [] - seen_paths = set() - for prim in stage.Traverse(): - if prim.GetTypeName() != "Camera": - continue - camera_path = str(prim.GetPath()) - binding = binding_by_camera_path.get(camera_path) - if binding is not None: - ordered_bindings.append(binding) - seen_paths.add(camera_path) - - for binding in ppisp_bindings: - if binding[0] not in seen_paths: - ordered_bindings.append(binding) - seen_paths.add(binding[0]) - return ordered_bindings - - -def format_available_ppisp_cameras(ppisp_bindings: list[tuple[str, Usd.Prim, Usd.Prim]]) -> str: - """Format PPISP-bound camera paths for CLI error messages.""" - return "\n ".join(dict.fromkeys(binding[0] for binding in ppisp_bindings)) - - -def resolve_source_camera_binding(source_stage: Usd.Stage) -> tuple[str, Usd.Prim, Usd.Prim]: - """Resolve the source camera and PPISP shader binding from CLI or source stage metadata.""" +def resolve_source_camera_binding(source_stage: Usd.Stage) -> tuple[str, Usd.Prim | None, Usd.Prim]: + """Resolve the source camera and PPISP camera binding from CLI or source stage metadata.""" ppisp_bindings = order_ppisp_bindings_by_camera(source_stage, find_ppisp_camera_bindings(source_stage)) if not ppisp_bindings: - raise RuntimeError("No PPISP RenderProduct bindings found in input scene.") + raise RuntimeError("No cameras with PPISP camera attributes found in input scene.") if args_cli.camera_prim_path is not None: camera_prim_path = args_cli.camera_prim_path @@ -260,14 +215,14 @@ def resolve_source_camera_binding(source_stage: Usd.Stage) -> tuple[str, Usd.Pri available = format_available_ppisp_cameras(ppisp_bindings) raise RuntimeError( f"Camera prim not found: {camera_prim_path}\n" - "Omit --camera_prim_path to auto-select a PPISP-bound camera, or use one of:\n" + "Omit --camera_prim_path to auto-select a camera with PPISP attributes, or use one of:\n" f" {available}" ) if camera_prim.GetTypeName() != "Camera": available = format_available_ppisp_cameras(ppisp_bindings) raise RuntimeError( f"Prim is not a Camera: {camera_prim_path} ({camera_prim.GetTypeName()})\n" - "Omit --camera_prim_path to auto-select a PPISP-bound camera, or use one of:\n" + "Omit --camera_prim_path to auto-select a camera with PPISP attributes, or use one of:\n" f" {available}" ) @@ -277,8 +232,8 @@ def resolve_source_camera_binding(source_stage: Usd.Stage) -> tuple[str, Usd.Pri available = format_available_ppisp_cameras(ppisp_bindings) raise RuntimeError( - f"Selected camera has no PPISP RenderProduct binding: {camera_prim_path}\n" - "Omit --camera_prim_path to auto-select a PPISP-bound camera, or use one of:\n" + f"Selected camera has no PPISP camera attributes: {camera_prim_path}\n" + "Omit --camera_prim_path to auto-select a camera with PPISP attributes, or use one of:\n" f" {available}" ) @@ -353,8 +308,10 @@ def bake_source_camera_pose_to_envs(source_stage: Usd.Stage, source_camera_prim_ ) -def get_render_product_resolution(render_product_prim: Usd.Prim) -> tuple[int, int] | None: +def get_render_product_resolution(render_product_prim: Usd.Prim | None) -> tuple[int, int] | None: """Return ``(width, height)`` from a RenderProduct ``resolution`` attribute.""" + if render_product_prim is None: + return None resolution_attr = render_product_prim.GetAttribute("resolution") if not resolution_attr: return None @@ -364,7 +321,7 @@ def get_render_product_resolution(render_product_prim: Usd.Prim) -> tuple[int, i return int(resolution[0]), int(resolution[1]) -def resolve_image_shape(render_product_prim: Usd.Prim) -> tuple[int, int]: +def resolve_image_shape(render_product_prim: Usd.Prim | None) -> tuple[int, int]: """Resolve demo output ``(width, height)`` preserving source aspect when height is omitted.""" width = args_cli.image_width height = args_cli.image_height @@ -380,14 +337,14 @@ def resolve_image_shape(render_product_prim: Usd.Prim) -> tuple[int, int]: return width, height -def make_ppisp_cfg(shader_prim: Usd.Prim, num_ppisp_bindings: int) -> PpispCfg: - """Parse the selected source PPISP shader into an explicit cfg for duplicated envs.""" - ppisp_cfg = ppisp_cfg_from_usd_shader(UsdShade.Shader(shader_prim)) - # The duplicated stage does not carry the source /Render scope; keep the - # parsed inputs as explicit values instead of resolving this source path later. - ppisp_cfg.shader_prim_path = None +def make_ppisp_cfg(camera_prim: Usd.Prim, num_ppisp_bindings: int) -> PpispCfg: + """Parse the selected source PPISP camera into an explicit cfg for duplicated envs.""" + ppisp_cfg = ppisp_cfg_from_usd_camera(camera_prim) + # The duplicated stage can remap source camera paths; keep the parsed inputs + # as explicit values instead of resolving the original camera path later. + ppisp_cfg.camera_prim_path = None if args_cli.ppisp_responsivity is None: - print(f"[INFO] Using USD-authored PPISP values from {num_ppisp_bindings} PPISP shader(s).", flush=True) + print(f"[INFO] Using USD-authored PPISP values from {num_ppisp_bindings} PPISP camera(s).", flush=True) else: ppisp_cfg.inputs["responsivity"] = float(args_cli.ppisp_responsivity) print( @@ -572,8 +529,8 @@ def main() -> None: source_stage = Usd.Stage.Open(args_cli.input_scene) if source_stage is None: raise RuntimeError(f"Failed to open input scene: {args_cli.input_scene}") - source_camera_prim_path, render_product_prim, shader_prim = resolve_source_camera_binding(source_stage) - ppisp_cfg = make_ppisp_cfg(shader_prim, len(find_ppisp_camera_bindings(source_stage))) + source_camera_prim_path, render_product_prim, ppisp_camera_prim = resolve_source_camera_binding(source_stage) + ppisp_cfg = make_ppisp_cfg(ppisp_camera_prim, len(find_ppisp_camera_bindings(source_stage))) camera_prim_path = source_camera_path_to_env_regex(source_stage, source_camera_prim_path) width, height = resolve_image_shape(render_product_prim) diff --git a/source/isaaclab/changelog.d/nicolasm-ppisp-controller-warp.skip b/source/isaaclab/changelog.d/nicolasm-ppisp-controller-warp.skip new file mode 100644 index 000000000000..1d58ec3c07e1 --- /dev/null +++ b/source/isaaclab/changelog.d/nicolasm-ppisp-controller-warp.skip @@ -0,0 +1 @@ +No user-facing isaaclab package changelog entry; PPISP controller coverage updates only. diff --git a/source/isaaclab/isaaclab/sensors/camera/camera_cfg.py b/source/isaaclab/isaaclab/sensors/camera/camera_cfg.py index 172756f7f95d..aaaa95c17283 100644 --- a/source/isaaclab/isaaclab/sensors/camera/camera_cfg.py +++ b/source/isaaclab/isaaclab/sensors/camera/camera_cfg.py @@ -203,8 +203,9 @@ class OffsetCfg: * ``None`` — ISP disabled. No HDR AOV is requested and no RTX-side tonemapping flags are flipped. - * A :class:`CameraISPMode` sentinel — the renderer backend walks the USD stage to - discover an ISP shader (e.g. via the :mod:`isaaclab_ppisp` package). + * A :class:`CameraISPMode` sentinel — the renderer backend discovers + camera-authored ``ppisp:*`` attributes (e.g. via the :mod:`isaaclab_ppisp` + package). * A concrete ISP cfg dataclass (e.g. :class:`isaaclab_ppisp.PpispCfg`) — used directly. The cfg applies once per Camera sensor batch. The PPISP Warp kernel takes diff --git a/source/isaaclab/isaaclab/sensors/camera/camera_isp.py b/source/isaaclab/isaaclab/sensors/camera/camera_isp.py index 8934198bbcf8..9fe61aedd209 100644 --- a/source/isaaclab/isaaclab/sensors/camera/camera_isp.py +++ b/source/isaaclab/isaaclab/sensors/camera/camera_isp.py @@ -29,16 +29,16 @@ class CameraISPMode(StrEnum): AUTO_CAMERA = "auto_camera" """Discover an ISP cfg bound to the batch's first matched camera prim. - The renderer walks the stage for a ``RenderProduct`` whose ``camera`` - relationship targets the first matched camera prim in the sensor batch. If - found and it has an ISP shader child, parses the shader. Resolves to - ``None`` (no ISP) if no match. + The renderer checks the first matched camera prim in the sensor batch for + recognized ``ppisp:*`` attributes, including optional controller weights. + It parses those attributes into an ISP cfg, or resolves to ``None`` when + the camera has no PPISP attributes. """ AUTO_ANY = "auto_any" - """Same as :attr:`AUTO_CAMERA`, then fall back to the first ISP shader anywhere on the stage. + """Same as :attr:`AUTO_CAMERA`, then fall back to the first PPISP camera anywhere on the stage. - Used to honour a stage-wide ISP authoring even when no ``RenderProduct`` - binds the shader explicitly to the batch's first matched camera prim. - Resolves to ``None`` only when the stage contains no ISP shader at all. + Used to honor stage-wide camera authoring when the batch's first matched + camera has no recognized ``ppisp:*`` attributes. Resolves to ``None`` only + when the stage contains no camera with PPISP attributes. """ diff --git a/source/isaaclab/test/sensors/generate_synthetic_gaussian_asset.py b/source/isaaclab/test/sensors/generate_synthetic_gaussian_asset.py index 2e6c9f471b04..daea8287b9a6 100644 --- a/source/isaaclab/test/sensors/generate_synthetic_gaussian_asset.py +++ b/source/isaaclab/test/sensors/generate_synthetic_gaussian_asset.py @@ -26,12 +26,13 @@ import torch from isaaclab_ppisp import PpispCfg, normalize_ppisp_cfg -from pxr import Gf, Sdf, Usd, UsdGeom +from pxr import Gf, Sdf, Usd, UsdGeom, Vt import isaaclab.sim as sim_utils from isaaclab.assets import AssetBaseCfg, RigidObjectCfg from isaaclab.scene import InteractiveScene, InteractiveSceneCfg from isaaclab.sensors.camera import Camera, CameraCfg +from isaaclab.sensors.camera.camera_isp import CameraISPMode from isaaclab.terrains import TerrainImporterCfg from isaaclab.utils.configclass import configclass @@ -230,6 +231,32 @@ def _attr(name: str, type_name: Sdf.ValueTypeName, value): _AGGRESSIVE_EXPOSURE_OFFSET = -5.0 +_PPISP_CONTROLLER_EXPECTED_WEIGHTS_LEN = 241_961 +_PPISP_CONTROLLER_OFF_CONV1_W = 0 +_PPISP_CONTROLLER_OFF_CONV1_B = _PPISP_CONTROLLER_OFF_CONV1_W + 16 * 3 +_PPISP_CONTROLLER_OFF_CONV2_W = _PPISP_CONTROLLER_OFF_CONV1_B + 16 +_PPISP_CONTROLLER_OFF_CONV2_B = _PPISP_CONTROLLER_OFF_CONV2_W + 32 * 16 +_PPISP_CONTROLLER_OFF_CONV3_W = _PPISP_CONTROLLER_OFF_CONV2_B + 32 +_PPISP_CONTROLLER_OFF_CONV3_B = _PPISP_CONTROLLER_OFF_CONV3_W + 64 * 32 +_PPISP_CONTROLLER_OFF_TRUNK0_W = _PPISP_CONTROLLER_OFF_CONV3_B + 64 +_PPISP_CONTROLLER_OFF_TRUNK0_B = _PPISP_CONTROLLER_OFF_TRUNK0_W + 128 * 1601 +_PPISP_CONTROLLER_OFF_TRUNK1_W = _PPISP_CONTROLLER_OFF_TRUNK0_B + 128 +_PPISP_CONTROLLER_OFF_TRUNK1_B = _PPISP_CONTROLLER_OFF_TRUNK1_W + 128 * 128 +_PPISP_CONTROLLER_OFF_TRUNK2_W = _PPISP_CONTROLLER_OFF_TRUNK1_B + 128 +_PPISP_CONTROLLER_OFF_TRUNK2_B = _PPISP_CONTROLLER_OFF_TRUNK2_W + 128 * 128 +_PPISP_CONTROLLER_OFF_EXP_W = _PPISP_CONTROLLER_OFF_TRUNK2_B + 128 +_PPISP_CONTROLLER_OFF_EXP_B = _PPISP_CONTROLLER_OFF_EXP_W + 128 +_PPISP_CONTROLLER_OFF_COL_W = _PPISP_CONTROLLER_OFF_EXP_B + 1 +_PPISP_CONTROLLER_OFF_COL_B = _PPISP_CONTROLLER_OFF_COL_W + 8 * 128 + +_PPISP_CONTROLLER_TOTAL_WEIGHTS = _PPISP_CONTROLLER_OFF_COL_B + 8 +if _PPISP_CONTROLLER_TOTAL_WEIGHTS != _PPISP_CONTROLLER_EXPECTED_WEIGHTS_LEN: + raise RuntimeError( + "Synthetic PPISP controller fixture offsets are inconsistent: " + f"{_PPISP_CONTROLLER_TOTAL_WEIGHTS} != {_PPISP_CONTROLLER_EXPECTED_WEIGHTS_LEN}." + ) + + def make_aggressive_ppisp_cfg(*, responsivity: float = 1.0) -> PpispCfg: """Return a :class:`~isaaclab_ppisp.PpispCfg` with every PPISP feature engaged enough to be assertable in a downstream test. @@ -301,6 +328,105 @@ def make_aggressive_ppisp_cfg(*, responsivity: float = 1.0) -> PpispCfg: return normalize_ppisp_cfg(PpispCfg(inputs=inputs)) +def make_neutral_ppisp_cfg(*, responsivity: float = 1.0) -> PpispCfg: + """Return a mild static PPISP cfg used as the camera-attribute negative control.""" + return normalize_ppisp_cfg(PpispCfg(inputs={"responsivity": responsivity, "exposureOffset": 0.0})) + + +def assert_images_meaningfully_different( + reference_rgb: torch.Tensor, + candidate_rgb: torch.Tensor, + *, + min_mean_abs_diff: float = 3.0, + label: str = "", +) -> None: + """Assert two LDR RGB tiles differ enough to prove PPISP attributes changed output.""" + prefix = f"[{label}] " if label else "" + diff = (reference_rgb[..., :3].float() - candidate_rgb[..., :3].float()).abs() + mean_abs_diff = diff.mean().item() + assert mean_abs_diff > min_mean_abs_diff, ( + f"{prefix}image difference too small: mean_abs_diff={mean_abs_diff:.3f}, " + f"expected > {min_mean_abs_diff}. The authored PPISP camera attributes may not be applied." + ) + + +def assert_ppisp_controller_matches_static( + static_rgb: torch.Tensor, + controller_rgb: torch.Tensor, + *, + max_mean_abs_diff: float = 8.0, + label: str = "", +) -> None: + """Assert deterministic controller output matches the equivalent static PPISP cfg.""" + prefix = f"[{label}] " if label else "" + diff = (static_rgb[..., :3].float() - controller_rgb[..., :3].float()).abs() + mean_abs_diff = diff.mean().item() + assert mean_abs_diff < max_mean_abs_diff, ( + f"{prefix}controller PPISP differs from static reference: mean_abs_diff={mean_abs_diff:.3f}, " + f"expected < {max_mean_abs_diff}." + ) + + +def _deterministic_controller_weights(ppisp_cfg: PpispCfg) -> tuple[float, ...]: + inputs = ppisp_cfg.inputs + weights = [0.0] * _PPISP_CONTROLLER_EXPECTED_WEIGHTS_LEN + weights[_PPISP_CONTROLLER_OFF_EXP_B] = float(inputs["exposureOffset"]) + color_values = ( + *_float2(inputs["colorLatentBlue"]), + *_float2(inputs["colorLatentRed"]), + *_float2(inputs["colorLatentGreen"]), + *_float2(inputs["colorLatentNeutral"]), + ) + for i, value in enumerate(color_values): + weights[_PPISP_CONTROLLER_OFF_COL_B + i] = value + return tuple(weights) + + +def _float2(value: float | tuple[float, float]) -> tuple[float, float]: + assert not isinstance(value, float) + return (float(value[0]), float(value[1])) + + +def _camera_path_for_env(env_id: int = 0) -> str: + return f"/World/envs/env_{env_id}/{SYNTHETIC_GAUSSIAN_SCENE_REL_PATH}/Cameras/{SYNTHETIC_GAUSSIAN_CAMERA_NAME}" + + +def _set_ppisp_camera_attrs( + stage: Usd.Stage, + inputs: dict[str, float | tuple[float, float]], + *, + controller_weights: tuple[float, ...] | None = None, +) -> None: + camera_prim = stage.GetPrimAtPath(_camera_path_for_env(0)) + if not camera_prim or not camera_prim.IsValid(): + raise RuntimeError(f"Synthetic PPISP camera prim not found: {_camera_path_for_env(0)}") + for name, value in inputs.items(): + if isinstance(value, tuple): + camera_prim.CreateAttribute(f"ppisp:{name}", Sdf.ValueTypeNames.Float2).Set( + Gf.Vec2f(float(value[0]), float(value[1])) + ) + else: + camera_prim.CreateAttribute(f"ppisp:{name}", Sdf.ValueTypeNames.Float).Set(float(value)) + if controller_weights is not None: + camera_prim.CreateAttribute("ppisp:controllerWeights", Sdf.ValueTypeNames.FloatArray).Set( + Vt.FloatArray(controller_weights) + ) + + +def author_static_ppisp_camera_attrs(stage: Usd.Stage, *, ppisp_cfg: PpispCfg) -> None: + """Author static PPISP camera attributes on the synthetic camera.""" + _set_ppisp_camera_attrs(stage, ppisp_cfg.inputs) + + +def author_controller_ppisp_camera_attrs(stage: Usd.Stage, *, ppisp_cfg: PpispCfg) -> None: + """Author PPISP camera attributes plus deterministic controller weights.""" + _set_ppisp_camera_attrs( + stage, + ppisp_cfg.inputs, + controller_weights=_deterministic_controller_weights(ppisp_cfg), + ) + + def assert_ppisp_invariants( rgb_tile: torch.Tensor, *, @@ -563,3 +689,95 @@ def render_synthetic_gaussian_scene( outputs = {name: tensor.clone().detach().cpu().to(torch.float32) for name, tensor in camera.data.output.items()} del camera return outputs + + +def render_synthetic_gaussian_scene_with_static_ppisp_attrs( + usd_path: str, + *, + sim_cfg: SimulationCfg, + renderer_cfg: RendererCfg, + ppisp_cfg: PpispCfg, + data_types: list[str], + num_envs: int = 1, + height: int = 128, + width: int = 128, + sim_dt: float = 1.0 / 60.0, + stabilisation_steps: int = 5, +) -> dict[str, torch.Tensor]: + """Render the synthesised gaussian asset through authored static PPISP camera attributes. + + The camera uses :class:`CameraISPMode.AUTO_CAMERA`; renderer backends must + discover the camera-authored PPISP attributes and route them through their + PPISP workflow. + """ + with fresh_synthetic_gaussian_interactive_scene(usd_path, sim_cfg, num_envs=num_envs) as sim: + author_static_ppisp_camera_attrs(sim.stage, ppisp_cfg=ppisp_cfg) + return _render_synthetic_gaussian_camera( + renderer_cfg=renderer_cfg, + data_types=data_types, + height=height, + width=width, + sim_dt=sim_dt, + stabilisation_steps=stabilisation_steps, + isp_cfg=CameraISPMode.AUTO_CAMERA, + sim=sim, + ) + + +def render_synthetic_gaussian_scene_with_controller_ppisp_attrs( + usd_path: str, + *, + sim_cfg: SimulationCfg, + renderer_cfg: RendererCfg, + ppisp_cfg: PpispCfg, + data_types: list[str], + num_envs: int = 1, + height: int = 128, + width: int = 128, + sim_dt: float = 1.0 / 60.0, + stabilisation_steps: int = 5, +) -> dict[str, torch.Tensor]: + """Render the synthesised gaussian asset through camera-authored controller weights.""" + with fresh_synthetic_gaussian_interactive_scene(usd_path, sim_cfg, num_envs=num_envs) as sim: + author_controller_ppisp_camera_attrs(sim.stage, ppisp_cfg=ppisp_cfg) + return _render_synthetic_gaussian_camera( + renderer_cfg=renderer_cfg, + data_types=data_types, + height=height, + width=width, + sim_dt=sim_dt, + stabilisation_steps=stabilisation_steps, + isp_cfg=CameraISPMode.AUTO_CAMERA, + sim=sim, + ) + + +def _render_synthetic_gaussian_camera( + *, + renderer_cfg: RendererCfg, + data_types: list[str], + height: int, + width: int, + sim_dt: float, + stabilisation_steps: int, + isp_cfg: PpispCfg | CameraISPMode | None, + sim: SimulationContext, +) -> dict[str, torch.Tensor]: + cfg = CameraCfg( + prim_path=SYNTHETIC_GAUSSIAN_CAMERA_REGEX, + update_period=0.0, + height=height, + width=width, + data_types=data_types, + spawn=None, + isp_cfg=isp_cfg, + renderer_cfg=renderer_cfg, + ) + camera = Camera(cfg) + sim.reset() + for _ in range(stabilisation_steps): + sim.step() + camera.update(sim_dt) + outputs = {name: tensor.clone().detach().cpu().to(torch.float32) for name, tensor in camera.data.output.items()} + del camera + return outputs diff --git a/source/isaaclab/test/sensors/test_camera_ppisp_gaussian.py b/source/isaaclab/test/sensors/test_camera_ppisp_gaussian.py index e143251f934a..6f1e681efec9 100644 --- a/source/isaaclab/test/sensors/test_camera_ppisp_gaussian.py +++ b/source/isaaclab/test/sensors/test_camera_ppisp_gaussian.py @@ -42,10 +42,16 @@ import pytest from generate_synthetic_gaussian_asset import ( SYNTHETIC_GAUSSIAN_CAMERA_REGEX, + assert_images_meaningfully_different, + assert_ppisp_controller_matches_static, assert_ppisp_invariants, assert_ppisp_lifts_exposure, + make_aggressive_ppisp_cfg, + make_neutral_ppisp_cfg, make_synthetic_gaussian_usd, render_synthetic_gaussian_scene, + render_synthetic_gaussian_scene_with_controller_ppisp_attrs, + render_synthetic_gaussian_scene_with_static_ppisp_attrs, ) from isaaclab.sim import SimulationCfg @@ -75,6 +81,13 @@ def _collect_renderer_cfg_params() -> list: SIM_DT = 0.01 MULTI_TILE_COUNT = 4 ISAAC_RTX_RESPONSIVITY = 1.2 +ISAAC_RTX_CONTROLLER_COMPARE_RESOLUTION = 640 +"""Resolution for static-vs-controller equivalence. + +Isaac RTX's very small render products go through DLSS upscaling paths whose +run-to-run variation can be close to the controller equivalence tolerance. +Keep only this comparison above that tiny-resolution regime. +""" def _isaac_rtx_sim_cfg(device: str) -> SimulationCfg: @@ -118,6 +131,71 @@ def test_camera_ppisp_wrapper_signatures_on_synthetic_gaussians(renderer_cfg_cls assert_ppisp_invariants(output["rgb"][0], label="isaac_rtx") +@pytest.mark.parametrize("device", ["cuda:0"]) +@pytest.mark.parametrize("renderer_cfg_cls", _RENDERER_CFG_PARAMS) +@pytest.mark.isaacsim_ci +def test_camera_ppisp_authored_static_attrs_are_applied_on_synthetic_gaussians(renderer_cfg_cls, device): + """Isaac RTX must apply camera-authored static PPISP attributes.""" + with tempfile.TemporaryDirectory(prefix="isaaclab-synth-gauss-") as tmpdir: + asset_path = make_synthetic_gaussian_usd(f"{tmpdir}/synthetic_gaussians.usda") + aggressive_cfg = make_aggressive_ppisp_cfg(responsivity=ISAAC_RTX_RESPONSIVITY) + + neutral = render_synthetic_gaussian_scene_with_static_ppisp_attrs( + asset_path, + sim_cfg=_isaac_rtx_sim_cfg(device), + renderer_cfg=renderer_cfg_cls(), + ppisp_cfg=make_neutral_ppisp_cfg(responsivity=ISAAC_RTX_RESPONSIVITY), + data_types=["rgb", "rgb_hdr"], + sim_dt=SIM_DT, + ) + aggressive = render_synthetic_gaussian_scene_with_static_ppisp_attrs( + asset_path, + sim_cfg=_isaac_rtx_sim_cfg(device), + renderer_cfg=renderer_cfg_cls(), + ppisp_cfg=aggressive_cfg, + data_types=["rgb", "rgb_hdr"], + sim_dt=SIM_DT, + ) + + assert_images_meaningfully_different(neutral["rgb"][0], aggressive["rgb"][0], label="isaac_rtx authored PPISP") + assert_ppisp_lifts_exposure(aggressive["rgb_hdr"][0], aggressive["rgb"][0], label="isaac_rtx authored PPISP") + assert_ppisp_invariants(aggressive["rgb"][0], label="isaac_rtx authored PPISP") + + +@pytest.mark.parametrize("device", ["cuda:0"]) +@pytest.mark.parametrize("renderer_cfg_cls", _RENDERER_CFG_PARAMS) +@pytest.mark.isaacsim_ci +def test_camera_ppisp_controller_matches_static_attrs_on_synthetic_gaussians(renderer_cfg_cls, device): + """Isaac RTX controller output must match the equivalent static PPISP cfg.""" + with tempfile.TemporaryDirectory(prefix="isaaclab-synth-gauss-") as tmpdir: + asset_path = make_synthetic_gaussian_usd(f"{tmpdir}/synthetic_gaussians.usda") + ppisp_cfg = make_aggressive_ppisp_cfg(responsivity=ISAAC_RTX_RESPONSIVITY) + + static = render_synthetic_gaussian_scene_with_static_ppisp_attrs( + asset_path, + sim_cfg=_isaac_rtx_sim_cfg(device), + renderer_cfg=renderer_cfg_cls(), + ppisp_cfg=ppisp_cfg, + data_types=["rgb", "rgb_hdr"], + height=ISAAC_RTX_CONTROLLER_COMPARE_RESOLUTION, + width=ISAAC_RTX_CONTROLLER_COMPARE_RESOLUTION, + sim_dt=SIM_DT, + ) + controller = render_synthetic_gaussian_scene_with_controller_ppisp_attrs( + asset_path, + sim_cfg=_isaac_rtx_sim_cfg(device), + renderer_cfg=renderer_cfg_cls(), + ppisp_cfg=ppisp_cfg, + data_types=["rgb", "rgb_hdr"], + height=ISAAC_RTX_CONTROLLER_COMPARE_RESOLUTION, + width=ISAAC_RTX_CONTROLLER_COMPARE_RESOLUTION, + sim_dt=SIM_DT, + ) + + assert_ppisp_controller_matches_static(static["rgb"][0], controller["rgb"][0], label="isaac_rtx controller") + assert_ppisp_invariants(controller["rgb"][0], label="isaac_rtx controller") + + @pytest.mark.parametrize("device", ["cuda:0"]) @pytest.mark.parametrize("renderer_cfg_cls", _RENDERER_CFG_PARAMS) @pytest.mark.isaacsim_ci diff --git a/source/isaaclab/test/sensors/test_camera_ppisp_gaussian_newton.py b/source/isaaclab/test/sensors/test_camera_ppisp_gaussian_newton.py index 809b67b42271..26aa6787e45a 100644 --- a/source/isaaclab/test/sensors/test_camera_ppisp_gaussian_newton.py +++ b/source/isaaclab/test/sensors/test_camera_ppisp_gaussian_newton.py @@ -52,10 +52,14 @@ from generate_synthetic_gaussian_asset import ( # noqa: E402 SYNTHETIC_GAUSSIAN_CAMERA_REGEX, + assert_ppisp_controller_matches_static, assert_ppisp_invariants, assert_ppisp_lifts_exposure, + make_aggressive_ppisp_cfg, make_synthetic_gaussian_usd, render_synthetic_gaussian_scene, + render_synthetic_gaussian_scene_with_controller_ppisp_attrs, + render_synthetic_gaussian_scene_with_static_ppisp_attrs, ) from isaaclab_newton.physics.mjwarp_manager_cfg import MJWarpSolverCfg # noqa: E402 from isaaclab_newton.physics.newton_manager_cfg import NewtonCfg # noqa: E402 @@ -65,10 +69,14 @@ else: tempfile = None SYNTHETIC_GAUSSIAN_CAMERA_REGEX = None + assert_ppisp_controller_matches_static = None assert_ppisp_invariants = None assert_ppisp_lifts_exposure = None + make_aggressive_ppisp_cfg = None make_synthetic_gaussian_usd = None render_synthetic_gaussian_scene = None + render_synthetic_gaussian_scene_with_controller_ppisp_attrs = None + render_synthetic_gaussian_scene_with_static_ppisp_attrs = None SimulationCfg = None MJWarpSolverCfg = None NewtonCfg = None @@ -114,6 +122,36 @@ def test_camera_ppisp_wrapper_signatures_on_synthetic_gaussians_newton(device): assert_ppisp_invariants(output["rgb"][0], label="newton_warp") +@pytest.mark.parametrize("device", ["cuda:0"]) +@pytest.mark.isaacsim_ci +@_SKIP_MISSING_NEWTON +def test_camera_ppisp_controller_matches_static_attrs_on_synthetic_gaussians_newton(device): + """Newton Warp renderer controller output must match the equivalent static PPISP cfg.""" + with tempfile.TemporaryDirectory(prefix="isaaclab-synth-gauss-") as tmpdir: + asset_path = make_synthetic_gaussian_usd(f"{tmpdir}/synthetic_gaussians.usda") + ppisp_cfg = make_aggressive_ppisp_cfg(responsivity=50.0) + + static = render_synthetic_gaussian_scene_with_static_ppisp_attrs( + asset_path, + sim_cfg=_newton_sim_cfg(device), + renderer_cfg=NewtonWarpRendererCfg(), + ppisp_cfg=ppisp_cfg, + data_types=["rgb", "rgb_hdr"], + sim_dt=SIM_DT, + ) + controller = render_synthetic_gaussian_scene_with_controller_ppisp_attrs( + asset_path, + sim_cfg=_newton_sim_cfg(device), + renderer_cfg=NewtonWarpRendererCfg(), + ppisp_cfg=ppisp_cfg, + data_types=["rgb", "rgb_hdr"], + sim_dt=SIM_DT, + ) + + assert_ppisp_controller_matches_static(static["rgb"][0], controller["rgb"][0], label="newton_warp controller") + assert_ppisp_invariants(controller["rgb"][0], label="newton_warp controller") + + @pytest.mark.parametrize("device", ["cuda:0"]) @pytest.mark.isaacsim_ci @_SKIP_MISSING_NEWTON diff --git a/source/isaaclab/test/sensors/test_camera_ppisp_gaussian_ovrtx.py b/source/isaaclab/test/sensors/test_camera_ppisp_gaussian_ovrtx.py index ca138fdedc27..642973c92f3d 100644 --- a/source/isaaclab/test/sensors/test_camera_ppisp_gaussian_ovrtx.py +++ b/source/isaaclab/test/sensors/test_camera_ppisp_gaussian_ovrtx.py @@ -49,10 +49,16 @@ import pytest from generate_synthetic_gaussian_asset import ( SYNTHETIC_GAUSSIAN_CAMERA_REGEX, + assert_images_meaningfully_different, + assert_ppisp_controller_matches_static, assert_ppisp_invariants, assert_ppisp_lifts_exposure, + make_aggressive_ppisp_cfg, + make_neutral_ppisp_cfg, make_synthetic_gaussian_usd, render_synthetic_gaussian_scene, + render_synthetic_gaussian_scene_with_controller_ppisp_attrs, + render_synthetic_gaussian_scene_with_static_ppisp_attrs, ) from isaaclab.sim import SimulationCfg @@ -130,6 +136,69 @@ def test_camera_ppisp_wrapper_signatures_on_synthetic_gaussians_ovrtx(device): assert_ppisp_invariants(output["rgb"][0], label="ovrtx") +@pytest.mark.parametrize("device", ["cuda:0"]) +@pytest.mark.isaacsim_ci +@_SKIP_MISSING_OVRTX +@_XFAIL_OVRTX_GAUSSIAN_PPISP +def test_camera_ppisp_authored_static_attrs_are_applied_on_synthetic_gaussians_ovrtx(device): + """OVRTX must apply camera-authored static PPISP attributes.""" + with tempfile.TemporaryDirectory(prefix="isaaclab-synth-gauss-") as tmpdir: + asset_path = make_synthetic_gaussian_usd(f"{tmpdir}/synthetic_gaussians.usda") + aggressive_cfg = make_aggressive_ppisp_cfg() + + neutral = render_synthetic_gaussian_scene_with_static_ppisp_attrs( + asset_path, + sim_cfg=_ovrtx_sim_cfg(device), + renderer_cfg=OVRTXRendererCfg(), + ppisp_cfg=make_neutral_ppisp_cfg(), + data_types=["rgb", "rgb_hdr"], + sim_dt=SIM_DT, + ) + aggressive = render_synthetic_gaussian_scene_with_static_ppisp_attrs( + asset_path, + sim_cfg=_ovrtx_sim_cfg(device), + renderer_cfg=OVRTXRendererCfg(), + ppisp_cfg=aggressive_cfg, + data_types=["rgb", "rgb_hdr"], + sim_dt=SIM_DT, + ) + + assert_images_meaningfully_different(neutral["rgb"][0], aggressive["rgb"][0], label="ovrtx authored PPISP") + assert_ppisp_lifts_exposure(aggressive["rgb_hdr"][0], aggressive["rgb"][0], label="ovrtx authored PPISP") + assert_ppisp_invariants(aggressive["rgb"][0], label="ovrtx authored PPISP") + + +@pytest.mark.parametrize("device", ["cuda:0"]) +@pytest.mark.isaacsim_ci +@_SKIP_MISSING_OVRTX +@_XFAIL_OVRTX_GAUSSIAN_PPISP +def test_camera_ppisp_controller_matches_static_attrs_on_synthetic_gaussians_ovrtx(device): + """OVRTX controller output must match the equivalent static PPISP cfg.""" + with tempfile.TemporaryDirectory(prefix="isaaclab-synth-gauss-") as tmpdir: + asset_path = make_synthetic_gaussian_usd(f"{tmpdir}/synthetic_gaussians.usda") + ppisp_cfg = make_aggressive_ppisp_cfg() + + static = render_synthetic_gaussian_scene_with_static_ppisp_attrs( + asset_path, + sim_cfg=_ovrtx_sim_cfg(device), + renderer_cfg=OVRTXRendererCfg(), + ppisp_cfg=ppisp_cfg, + data_types=["rgb", "rgb_hdr"], + sim_dt=SIM_DT, + ) + controller = render_synthetic_gaussian_scene_with_controller_ppisp_attrs( + asset_path, + sim_cfg=_ovrtx_sim_cfg(device), + renderer_cfg=OVRTXRendererCfg(), + ppisp_cfg=ppisp_cfg, + data_types=["rgb", "rgb_hdr"], + sim_dt=SIM_DT, + ) + + assert_ppisp_controller_matches_static(static["rgb"][0], controller["rgb"][0], label="ovrtx controller") + assert_ppisp_invariants(controller["rgb"][0], label="ovrtx controller") + + @pytest.mark.parametrize("device", ["cuda:0"]) @pytest.mark.isaacsim_ci @_SKIP_MISSING_OVRTX diff --git a/source/isaaclab_newton/changelog.d/nicolasm-camera-ppisp-controller.rst b/source/isaaclab_newton/changelog.d/nicolasm-camera-ppisp-controller.rst new file mode 100644 index 000000000000..6773ca02f485 --- /dev/null +++ b/source/isaaclab_newton/changelog.d/nicolasm-camera-ppisp-controller.rst @@ -0,0 +1,4 @@ +Added +^^^^^ + +* Added Newton camera renderer support for PPISP controller parameters. diff --git a/source/isaaclab_newton/isaaclab_newton/renderers/newton_warp_renderer.py b/source/isaaclab_newton/isaaclab_newton/renderers/newton_warp_renderer.py index e9cf7d9ee083..d5ba5bad0eeb 100644 --- a/source/isaaclab_newton/isaaclab_newton/renderers/newton_warp_renderer.py +++ b/source/isaaclab_newton/isaaclab_newton/renderers/newton_warp_renderer.py @@ -276,14 +276,15 @@ def prepare_cameras(self, stage: Any, spec: CameraRenderSpec) -> None: owns the sentinel-resolution + cfg-normalization step. Newton has no USD-side overrides to author beyond this. """ - if spec.cfg.isp_cfg is None or not spec.camera_prim_paths: + if spec.cfg.isp_cfg is None: return try: from isaaclab_ppisp import resolve_and_normalize except ModuleNotFoundError as exc: _raise_missing_ppisp_error(exc) - spec.cfg.isp_cfg = resolve_and_normalize(spec.cfg.isp_cfg, stage, spec.camera_prim_paths[0]) + camera_prim_path = spec.camera_prim_paths[0] if spec.camera_prim_paths else None + spec.cfg.isp_cfg = resolve_and_normalize(spec.cfg.isp_cfg, stage, camera_prim_path) def prepare_stage(self, stage: Any, num_envs: int) -> None: """No-op for Newton Warp - uses Newton scene directly without stage export. diff --git a/source/isaaclab_ov/changelog.d/nicolasm-camera-ppisp-controller.rst b/source/isaaclab_ov/changelog.d/nicolasm-camera-ppisp-controller.rst new file mode 100644 index 000000000000..e9f9e31a7382 --- /dev/null +++ b/source/isaaclab_ov/changelog.d/nicolasm-camera-ppisp-controller.rst @@ -0,0 +1,4 @@ +Fixed +^^^^^ + +* Fixed OVRTX camera PPISP/HDR rendering to disable Gaussian skip-tonemapping before renderer initialization. diff --git a/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_renderer.py b/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_renderer.py index 8363a3c2bdfd..0c5e2b82e38b 100644 --- a/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_renderer.py +++ b/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_renderer.py @@ -240,15 +240,16 @@ def prepare_cameras(self, stage: Any, spec: CameraRenderSpec) -> None: the RTX exposure model OVRTX embeds does not compound on top of the ISP. Without an ISP, the camera prim's authored exposure is left alone. """ - if not spec.camera_prim_paths or spec.cfg.isp_cfg is None: + if spec.cfg.isp_cfg is None: return try: from isaaclab_ppisp import apply_rtx_exposure_overrides, resolve_and_normalize except ModuleNotFoundError as exc: _raise_missing_ppisp_error(exc) - spec.cfg.isp_cfg = resolve_and_normalize(spec.cfg.isp_cfg, stage, spec.camera_prim_paths[0]) - if spec.cfg.isp_cfg is None: + camera_prim_path = spec.camera_prim_paths[0] if spec.camera_prim_paths else None + spec.cfg.isp_cfg = resolve_and_normalize(spec.cfg.isp_cfg, stage, camera_prim_path) + if spec.cfg.isp_cfg is None or not spec.camera_prim_paths: return apply_rtx_exposure_overrides(stage, list(spec.camera_prim_paths)) diff --git a/source/isaaclab_physx/changelog.d/nicolasm-ppisp-controller-warp.skip b/source/isaaclab_physx/changelog.d/nicolasm-ppisp-controller-warp.skip new file mode 100644 index 000000000000..17fd22cdd327 --- /dev/null +++ b/source/isaaclab_physx/changelog.d/nicolasm-ppisp-controller-warp.skip @@ -0,0 +1 @@ +Internal-only change: drop the now-redundant ``stage`` argument when constructing ``PpispPipeline`` (``isp_cfg`` is already resolved via ``resolve_and_normalize``). No user-facing change. diff --git a/source/isaaclab_physx/isaaclab_physx/renderers/isaac_rtx_renderer.py b/source/isaaclab_physx/isaaclab_physx/renderers/isaac_rtx_renderer.py index 9f25566fecf7..e679266423ef 100644 --- a/source/isaaclab_physx/isaaclab_physx/renderers/isaac_rtx_renderer.py +++ b/source/isaaclab_physx/isaaclab_physx/renderers/isaac_rtx_renderer.py @@ -126,15 +126,16 @@ def prepare_cameras(self, stage: Any, spec: CameraRenderSpec) -> None: RTX's physical-camera exposure model does not compound on top of the ISP. Without an ISP, the camera prim's authored exposure is left alone. """ - if not spec.camera_prim_paths or spec.cfg.isp_cfg is None: + if spec.cfg.isp_cfg is None: return try: from isaaclab_ppisp import apply_rtx_exposure_overrides, resolve_and_normalize except ModuleNotFoundError as exc: _raise_missing_ppisp_error(exc) - spec.cfg.isp_cfg = resolve_and_normalize(spec.cfg.isp_cfg, stage, spec.camera_prim_paths[0]) - if spec.cfg.isp_cfg is None: + camera_prim_path = spec.camera_prim_paths[0] if spec.camera_prim_paths else None + spec.cfg.isp_cfg = resolve_and_normalize(spec.cfg.isp_cfg, stage, camera_prim_path) + if spec.cfg.isp_cfg is None or not spec.camera_prim_paths: return apply_rtx_exposure_overrides(stage, list(spec.camera_prim_paths)) @@ -331,7 +332,7 @@ def create_render_data(self, spec: CameraRenderSpec) -> IsaacRtxRenderData: except ModuleNotFoundError as exc: _raise_missing_ppisp_error(exc) - ppisp_pipeline = PpispPipeline(spec.cfg.isp_cfg, stage=stage) + ppisp_pipeline = PpispPipeline(spec.cfg.isp_cfg) return IsaacRtxRenderData( annotators=annotators, diff --git a/source/isaaclab_ppisp/changelog.d/nicolasm-ppisp-controller-warp.minor.rst b/source/isaaclab_ppisp/changelog.d/nicolasm-ppisp-controller-warp.minor.rst new file mode 100644 index 000000000000..955c446f00dd --- /dev/null +++ b/source/isaaclab_ppisp/changelog.d/nicolasm-ppisp-controller-warp.minor.rst @@ -0,0 +1,8 @@ +Added +^^^^^ + +* Added PPISP camera-attribute discovery via :class:`~isaaclab.sensors.camera.CameraISPMode.AUTO_CAMERA` and + :class:`~isaaclab.sensors.camera.CameraISPMode.AUTO_ANY`, embedded controller-weight parsing, and native + CUDA controller execution for renderer PPISP output. +* Updated PPISP discovery documentation to describe camera-authored ``ppisp:*`` attributes rather than + RenderProduct-bound shader lookup. diff --git a/source/isaaclab_ppisp/isaaclab_ppisp/__init__.pyi b/source/isaaclab_ppisp/isaaclab_ppisp/__init__.pyi index f5bc313692ab..2d69e60325e9 100644 --- a/source/isaaclab_ppisp/isaaclab_ppisp/__init__.pyi +++ b/source/isaaclab_ppisp/isaaclab_ppisp/__init__.pyi @@ -4,33 +4,37 @@ # SPDX-License-Identifier: BSD-3-Clause __all__ = [ + "PPISP_CONTROLLER_EXPECTED_WEIGHTS_LEN", "PPISP_DEFAULT_INPUTS", - "PPISP_SHADER_NAME", "PpispCfg", "PpispPipeline", "apply_ppisp_to_rgba", + "apply_ppisp_to_rgba_with_controller_params", "apply_rtx_exposure_overrides", "auto_any_ppisp_cfg", "auto_camera_ppisp_cfg", + "compute_ppisp_controller_params", "default_ppisp_inputs", + "has_ppisp_camera_attrs", "normalize_ppisp_cfg", - "ppisp_cfg_from_usd_shader", + "ppisp_cfg_from_usd_camera", "ppisp_cfg_from_usd_stage", "resolve_and_normalize", ] from .cfg import ( + PPISP_CONTROLLER_EXPECTED_WEIGHTS_LEN, PPISP_DEFAULT_INPUTS, - PPISP_SHADER_NAME, PpispCfg, auto_any_ppisp_cfg, auto_camera_ppisp_cfg, default_ppisp_inputs, + has_ppisp_camera_attrs, normalize_ppisp_cfg, - ppisp_cfg_from_usd_shader, + ppisp_cfg_from_usd_camera, ppisp_cfg_from_usd_stage, resolve_and_normalize, ) -from .kernels import apply_ppisp_to_rgba +from .kernels import apply_ppisp_to_rgba, apply_ppisp_to_rgba_with_controller_params, compute_ppisp_controller_params from .pipeline import PpispPipeline from .rtx_camera_overrides import apply_rtx_exposure_overrides diff --git a/source/isaaclab_ppisp/isaaclab_ppisp/_demo_utils.py b/source/isaaclab_ppisp/isaaclab_ppisp/_demo_utils.py new file mode 100644 index 000000000000..64d44f8fcab5 --- /dev/null +++ b/source/isaaclab_ppisp/isaaclab_ppisp/_demo_utils.py @@ -0,0 +1,68 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Shared PPISP camera-selection helpers for demo scripts.""" + +from __future__ import annotations + +from typing import Any + +from pxr import Sdf + +from .cfg import has_ppisp_camera_attrs + +PpispCameraBinding = tuple[str, Any | None, Any] + + +def find_ppisp_camera_bindings(stage: Any) -> list[PpispCameraBinding]: + """Return cameras with recognized PPISP attributes and optional render-product metadata.""" + bindings = [] + for prim in stage.Traverse(): + if not has_ppisp_camera_attrs(prim): + continue + camera_path = str(prim.GetPath()) + bindings.append((camera_path, find_render_product_for_camera(stage, camera_path), prim)) + return bindings + + +def find_render_product_for_camera(stage: Any, camera_prim_path: str) -> Any | None: + """Return the first RenderProduct targeting ``camera_prim_path``, if any.""" + target_path = Sdf.Path(camera_prim_path) + for prim in stage.Traverse(): + if prim.GetTypeName() != "RenderProduct": + continue + camera_rel = prim.GetRelationship("camera") + if camera_rel and target_path in camera_rel.GetTargets(): + return prim + return None + + +def order_ppisp_bindings_by_camera(stage: Any, ppisp_bindings: list[PpispCameraBinding]) -> list[PpispCameraBinding]: + """Return PPISP bindings ordered by source camera prim traversal.""" + binding_by_camera_path = {} + for binding in ppisp_bindings: + binding_by_camera_path.setdefault(binding[0], binding) + + ordered_bindings = [] + seen_paths = set() + for prim in stage.Traverse(): + if prim.GetTypeName() != "Camera": + continue + camera_path = str(prim.GetPath()) + binding = binding_by_camera_path.get(camera_path) + if binding is not None: + ordered_bindings.append(binding) + seen_paths.add(camera_path) + + for binding in ppisp_bindings: + if binding[0] not in seen_paths: + ordered_bindings.append(binding) + seen_paths.add(binding[0]) + return ordered_bindings + + +def format_available_ppisp_cameras(ppisp_bindings: list[PpispCameraBinding]) -> str: + """Format cameras with PPISP attributes for CLI error messages.""" + return "\n ".join(dict.fromkeys(binding[0] for binding in ppisp_bindings)) diff --git a/source/isaaclab_ppisp/isaaclab_ppisp/cfg.py b/source/isaaclab_ppisp/isaaclab_ppisp/cfg.py index c2f0fa43bc8e..29275d2672c1 100644 --- a/source/isaaclab_ppisp/isaaclab_ppisp/cfg.py +++ b/source/isaaclab_ppisp/isaaclab_ppisp/cfg.py @@ -3,7 +3,7 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""PPISP configuration and USD/shader parsing helpers. +"""PPISP configuration and USD parsing helpers. The implementation follows the physically plausible ISP model described in https://arxiv.org/abs/2601.18336. @@ -16,8 +16,19 @@ from isaaclab.utils.configclass import configclass -PPISP_SHADER_NAME = "PPISP" -"""Conventional prim name for a PPISP shader child under a ``RenderProduct``.""" +PPISP_ATTR_NAMESPACE = "ppisp:" +"""Namespace prefix for authoritative PPISP attributes authored on a USD camera.""" + +PPISP_CONTROLLER_WEIGHTS_CAMERA_ATTR = "controllerWeights" +"""Camera ``ppisp:*`` attribute name containing flattened controller weights.""" + +PPISP_CONTROLLER_EXPECTED_WEIGHTS_LEN = 241_961 +"""Flattened element count of the camera-authored controller weight array exported by NRE. + +This is a frozen architectural constant tied to the exported controller network +shape (see :mod:`isaaclab_ppisp.kernels` for the offset layout). USD parsing +and Warp execution validate against it and fail loudly on a mismatch. +""" PPISP_FLOAT2_INPUTS = { "vignettingCenterR", @@ -72,15 +83,17 @@ def default_ppisp_inputs() -> dict[str, float | tuple[float, float]]: class PpispCfg: """Configuration for PPISP post-processing. - PPISP inputs are static in IsaacLab. If imported from animated USD shader inputs, - the first authored time sample is used and later samples are ignored. + PPISP inputs are static in IsaacLab. NRE exports store the authoritative + values on a USD camera as ``ppisp:*`` attributes. If animated USD + attributes are imported, the first authored time sample is used and later + samples are ignored. """ - shader_prim_path: str | None = None - """Optional source USD shader prim path used to populate :attr:`inputs`.""" + camera_prim_path: str | None = None + """Optional USD camera prim path used to import PPISP camera attributes.""" inputs: dict[str, float | tuple[float, float]] = field(default_factory=default_ppisp_inputs) - """Flat PPISP shader input values keyed by USD input name. + """Flat PPISP values keyed by PPISP parameter name. Coordinate conventions for spatial inputs: @@ -94,6 +107,26 @@ class PpispCfg: frame the image corners sit at ``r^2 = 0.5``. """ + controller_prior_exposure: float = 0.0 + """Controller prior exposure [EV] used by the native controller path.""" + + controller_responsivity: float | None = None + """Controller feature-extraction responsivity [dimensionless]. + + When ``None``, the controller uses the static PPISP ``responsivity`` camera + attribute so feature extraction sees the same responsivity-scaled HDR + radiance as the image PPISP transform. + """ + + controller_weights: tuple[float, ...] | None = None + """Flattened controller weights. + + USD imports read these from the camera's ``ppisp:controllerWeights`` + attribute. When present, the native controller predicts ``exposureOffset`` + and the four color latents from the HDR image each frame. Static PPISP + inputs still provide responsivity, vignetting, and CRF. + """ + def normalize_ppisp_cfg( ppisp_cfg: PpispCfg | None, @@ -102,9 +135,9 @@ def normalize_ppisp_cfg( """Normalise a :class:`PpispCfg` for downstream consumption. * If ``ppisp_cfg`` is ``None``, returns ``None``. - * If ``ppisp_cfg.shader_prim_path`` is set and ``stage`` is supplied, - merges the shader-authored values with the cfg's explicit overrides - (see :func:`_merge_shader_inputs_with_cfg`). + * If ``ppisp_cfg.camera_prim_path`` is set, requires ``stage`` and merges + camera-authored USD values with the cfg's explicit overrides (see + :func:`_merge_camera_attrs_with_cfg`). * Otherwise validates ``ppisp_cfg.inputs`` and fills in defaults. """ if ppisp_cfg is None: @@ -112,39 +145,58 @@ def normalize_ppisp_cfg( if not isinstance(ppisp_cfg, PpispCfg): raise TypeError(f"Unsupported PPISP configuration type: {type(ppisp_cfg)!r}") input_overrides = dict(ppisp_cfg.inputs) - if ppisp_cfg.shader_prim_path and stage is not None: - return _merge_shader_inputs_with_cfg(ppisp_cfg, stage, input_overrides) + if ppisp_cfg.camera_prim_path: + if stage is None: + raise ValueError("PpispCfg.camera_prim_path requires a USD stage for normalization.") + return _merge_camera_attrs_with_cfg(ppisp_cfg, stage, input_overrides) ppisp_cfg.inputs = _normalized_inputs(input_overrides) + _finalize_ppisp_cfg(ppisp_cfg) return ppisp_cfg -def ppisp_cfg_from_usd_shader(shader: Any) -> PpispCfg: - """Create :class:`PpispCfg` from a ``UsdShade.Shader`` prim. +def ppisp_cfg_from_usd_camera(camera_prim: Any) -> PpispCfg: + """Create :class:`PpispCfg` from a USD camera prim. - Animated inputs are collapsed to their first authored time sample. + PPISP values are read from camera ``ppisp:*`` attributes. Animated + attributes are collapsed to their first authored time sample. """ - cfg = PpispCfg(shader_prim_path=str(shader.GetPath())) - values = default_ppisp_inputs() - for input_name in values: - shader_input = shader.GetInput(input_name) - if not shader_input: - continue - attr = shader_input.GetAttr() - value = _read_first_authored_value(attr) - if value is not None: - values[input_name] = _normalize_input_value(input_name, value) + cfg = _ppisp_cfg_from_usd_camera(camera_prim) + _finalize_ppisp_cfg(cfg) + cfg.camera_prim_path = None + return cfg + + +def _ppisp_cfg_from_usd_camera(camera_prim: Any) -> PpispCfg: + values = _read_ppisp_inputs_from_camera(camera_prim) + controller_weights = _read_controller_weights_from_camera(camera_prim) + if values is None and controller_weights is None: + camera_path = str(camera_prim.GetPath()) if camera_prim and camera_prim.IsValid() else "" + raise ValueError( + f"PPISP camera attributes were not found on camera {camera_path}; expected ppisp:* attributes." + ) + + cfg = PpispCfg(camera_prim_path=str(camera_prim.GetPath())) + if values is None: + values = default_ppisp_inputs() cfg.inputs = values + if controller_weights is not None: + cfg.controller_weights = controller_weights return cfg -def ppisp_cfg_from_usd_stage(stage: Any, shader_prim_path: str) -> PpispCfg: - """Create :class:`PpispCfg` from a shader prim path in a USD stage.""" - from pxr import UsdShade +def ppisp_cfg_from_usd_stage(stage: Any, camera_prim_path: str) -> PpispCfg: + """Create :class:`PpispCfg` from a camera prim path in a USD stage.""" - shader = UsdShade.Shader(stage.GetPrimAtPath(shader_prim_path)) - if not shader: - raise ValueError(f"PPISP shader prim not found at path: {shader_prim_path}") - return ppisp_cfg_from_usd_shader(shader) + return ppisp_cfg_from_usd_camera(_get_camera_prim_at_path(stage, camera_prim_path)) + + +def _get_camera_prim_at_path(stage: Any, camera_prim_path: str) -> Any: + camera_prim = stage.GetPrimAtPath(camera_prim_path) + if not camera_prim or not camera_prim.IsValid(): + raise ValueError(f"PPISP camera prim not found at path: {camera_prim_path}") + if camera_prim.GetTypeName() != "Camera": + raise ValueError(f"PPISP prim is not a Camera: {camera_prim_path} ({camera_prim.GetTypeName()})") + return camera_prim def _normalized_inputs(inputs: dict[str, Any]) -> dict[str, float | tuple[float, float]]: @@ -156,14 +208,25 @@ def _normalized_inputs(inputs: dict[str, Any]) -> dict[str, float | tuple[float, return values -def _merge_shader_inputs_with_cfg( +def _merge_camera_attrs_with_cfg( ppisp_cfg: PpispCfg, stage: Any, input_overrides: dict[str, Any], ) -> PpispCfg: - parsed_cfg = ppisp_cfg_from_usd_stage(stage, ppisp_cfg.shader_prim_path) - if input_overrides != PPISP_DEFAULT_INPUTS: - parsed_cfg.inputs.update(_normalized_input_overrides(input_overrides)) + assert ppisp_cfg.camera_prim_path is not None + parsed_cfg = _ppisp_cfg_from_usd_camera(_get_camera_prim_at_path(stage, ppisp_cfg.camera_prim_path)) + normalized_overrides = _normalized_input_overrides(input_overrides) + if normalized_overrides != PPISP_DEFAULT_INPUTS: + parsed_cfg.inputs.update(normalized_overrides) + if ppisp_cfg.controller_weights is not None: + parsed_cfg.controller_prior_exposure = ppisp_cfg.controller_prior_exposure + parsed_cfg.controller_weights = ppisp_cfg.controller_weights + elif ppisp_cfg.controller_prior_exposure != 0.0: + parsed_cfg.controller_prior_exposure = ppisp_cfg.controller_prior_exposure + if ppisp_cfg.controller_responsivity is not None: + parsed_cfg.controller_responsivity = ppisp_cfg.controller_responsivity + _finalize_ppisp_cfg(parsed_cfg) + parsed_cfg.camera_prim_path = None return parsed_cfg @@ -184,6 +247,13 @@ def _normalize_input_value(input_name: str, value: Any) -> float | tuple[float, return float(value) +def _finalize_ppisp_cfg(ppisp_cfg: PpispCfg) -> None: + if ppisp_cfg.controller_responsivity is None: + ppisp_cfg.controller_responsivity = float(ppisp_cfg.inputs["responsivity"]) + else: + ppisp_cfg.controller_responsivity = float(ppisp_cfg.controller_responsivity) + + def _read_first_authored_value(attr: Any) -> Any: time_samples = attr.GetTimeSamples() if time_samples: @@ -191,32 +261,93 @@ def _read_first_authored_value(attr: Any) -> Any: return attr.Get() -def resolve_and_normalize(isp_cfg: Any, stage: Any, camera_prim_path: str) -> PpispCfg | None: +def _read_ppisp_inputs_from_camera(camera_prim: Any | None) -> dict[str, float | tuple[float, float]] | None: + if camera_prim is None or not camera_prim.IsValid(): + return None + + values = default_ppisp_inputs() + found = False + for input_name in values: + attr = camera_prim.GetAttribute(f"{PPISP_ATTR_NAMESPACE}{input_name}") + if not attr or not attr.IsValid(): + continue + value = _read_first_authored_value(attr) + if value is not None: + values[input_name] = _normalize_input_value(input_name, value) + found = True + return values if found else None + + +def _read_controller_weights_from_camera(camera_prim: Any | None) -> tuple[float, ...] | None: + if camera_prim is None or not camera_prim.IsValid(): + return None + attr = camera_prim.GetAttribute(f"{PPISP_ATTR_NAMESPACE}{PPISP_CONTROLLER_WEIGHTS_CAMERA_ATTR}") + if not attr or not attr.IsValid(): + return None + value = _read_first_authored_value(attr) + if value is None: + return None + weights = tuple(float(v) for v in value) + if len(weights) != PPISP_CONTROLLER_EXPECTED_WEIGHTS_LEN: + raise ValueError( + "Expected " + f"{PPISP_CONTROLLER_EXPECTED_WEIGHTS_LEN} PPISP controller weights on camera " + f"{camera_prim.GetPath()}, got {len(weights)}." + ) + return weights + + +def _has_ppisp_camera_attrs(camera_prim: Any | None) -> bool: + if camera_prim is None or not camera_prim.IsValid() or camera_prim.GetTypeName() != "Camera": + return False + for input_name in PPISP_DEFAULT_INPUTS: + attr = camera_prim.GetAttribute(f"{PPISP_ATTR_NAMESPACE}{input_name}") + if attr and attr.IsValid() and _read_first_authored_value(attr) is not None: + return True + weights_attr = camera_prim.GetAttribute(f"{PPISP_ATTR_NAMESPACE}{PPISP_CONTROLLER_WEIGHTS_CAMERA_ATTR}") + return bool(weights_attr and weights_attr.IsValid() and _read_first_authored_value(weights_attr) is not None) + + +def has_ppisp_camera_attrs(camera_prim: Any | None) -> bool: + """Return whether a USD camera prim contains recognized PPISP camera attributes. + + Args: + camera_prim: USD prim to inspect. + + Returns: + True when ``camera_prim`` is a camera with at least one recognized + ``ppisp:*`` attribute, otherwise false. + """ + return _has_ppisp_camera_attrs(camera_prim) + + +def resolve_and_normalize(isp_cfg: Any, stage: Any, camera_prim_path: str | None = None) -> PpispCfg | None: """Resolve a Camera sensor batch's ``isp_cfg`` to a normalised cfg or ``None``. Handles all three legal forms of :attr:`~isaaclab.sensors.camera.CameraCfg.isp_cfg`: * ``None`` → returns ``None``. - * :class:`~isaaclab.sensors.camera.CameraISPMode` sentinel — walks the stage - via :func:`auto_camera_ppisp_cfg` (and :func:`auto_any_ppisp_cfg` for - ``AUTO_ANY``) to discover a PPISP shader. Returns the parsed + - normalised :class:`PpispCfg`, or ``None`` if no shader matched. + * :class:`~isaaclab.sensors.camera.CameraISPMode` sentinel — checks the + target camera via :func:`auto_camera_ppisp_cfg` (and uses + :func:`auto_any_ppisp_cfg` for ``AUTO_ANY`` or when no camera path is + supplied) to discover a PPISP camera. Returns the parsed + normalised + :class:`PpispCfg`, or ``None`` if no PPISP camera matched. * Concrete :class:`PpispCfg` — normalises in place (validates input keys, - fills defaults, and merges shader-authored values when ``shader_prim_path`` - is set). + fills defaults, and merges camera-authored USD values when + ``camera_prim_path`` is set). This is the single entry point renderer backends call inside their ``prepare_cameras`` hook so :mod:`isaaclab.sensors.camera` does not need to know about PPISP types at all. The returned cfg applies to the whole Camera sensor batch; callers pass the first matched camera prim path for - the camera-bound discovery phase. + the camera-local discovery phase. Args: isp_cfg: The Camera sensor's :attr:`isp_cfg` value (``None``, ``CameraISPMode``, or :class:`PpispCfg`). - stage: USD stage used for sentinel discovery and shader-path resolution. - camera_prim_path: Absolute path of the first matched camera prim in the - Camera sensor batch (target of the ``camera`` relationship for the - camera-bound discovery phase). + stage: USD stage used for sentinel discovery and camera-path resolution. + camera_prim_path: Optional absolute path of the first matched camera + prim in the Camera sensor batch. When omitted, discovery uses the + first camera on the stage with PPISP camera attributes. Returns: A fully-normalised :class:`PpispCfg`, or ``None`` if the batch has no ISP. @@ -227,73 +358,51 @@ def resolve_and_normalize(isp_cfg: Any, stage: Any, camera_prim_path: str) -> Pp if isp_cfg is None: return None if isinstance(isp_cfg, CameraISPMode): - resolved = auto_camera_ppisp_cfg(stage, camera_prim_path) - if resolved is None and isp_cfg == CameraISPMode.AUTO_ANY: + resolved = auto_camera_ppisp_cfg(stage, camera_prim_path) if camera_prim_path else None + if resolved is None and (isp_cfg == CameraISPMode.AUTO_ANY or not camera_prim_path): resolved = auto_any_ppisp_cfg(stage) if resolved is None: return None - return normalize_ppisp_cfg(resolved, stage=stage) + return normalize_ppisp_cfg(resolved) return normalize_ppisp_cfg(isp_cfg, stage=stage) def auto_camera_ppisp_cfg(stage: Any, camera_prim_path: str) -> PpispCfg | None: - """Find the first PPISP shader on ``stage`` bound to ``camera_prim_path``. + """Find PPISP camera attributes for ``camera_prim_path`` on ``stage``. - Walks ``stage`` looking for the first ``RenderProduct`` whose ``camera`` - relationship targets ``camera_prim_path``. If that ``RenderProduct`` has a - child shader at ``/``, parses - and returns the corresponding :class:`PpispCfg`. Returns ``None`` if - no matching ``RenderProduct`` is found, or if it has no PPISP shader child. + Checks only the target camera itself. Use :func:`auto_any_ppisp_cfg` when + the caller intentionally wants a stage-wide fallback. Args: stage: USD stage to search. - camera_prim_path: Absolute camera prim path the ``RenderProduct``'s - ``camera`` relationship must target. + camera_prim_path: Absolute camera prim path to match. Returns: - Parsed :class:`PpispCfg` if a matching shader was found, else ``None``. + Parsed :class:`PpispCfg` if a matching camera was found, else ``None``. """ - from pxr import Sdf, UsdShade - - target_path = Sdf.Path(camera_prim_path) - for prim in stage.Traverse(): - if prim.GetTypeName() != "RenderProduct": - continue - camera_rel = prim.GetRelationship("camera") - if not camera_rel: - continue - if target_path not in camera_rel.GetTargets(): - continue - shader_path = prim.GetPath().AppendChild(PPISP_SHADER_NAME) - shader_prim = stage.GetPrimAtPath(shader_path) - if shader_prim and shader_prim.IsValid(): - return ppisp_cfg_from_usd_shader(UsdShade.Shader(shader_prim)) - return None + camera_prim = stage.GetPrimAtPath(camera_prim_path) + if not _has_ppisp_camera_attrs(camera_prim): + return None + cfg = ppisp_cfg_from_usd_camera(camera_prim) + cfg.camera_prim_path = None + return cfg def auto_any_ppisp_cfg(stage: Any) -> PpispCfg | None: - """Find the first PPISP shader anywhere on ``stage``, regardless of binding. + """Find the first camera with PPISP attributes anywhere on ``stage``. - Walks ``stage`` looking for the first prim named - :data:`PPISP_SHADER_NAME` that resolves to a valid ``UsdShade.Shader``. - Returns ``None`` if no such shader exists. - - Used as a fallback when no ``RenderProduct`` binds a PPISP shader to a - given camera but the scene still contains a PPISP configuration that - should apply. + Used as a fallback when no camera is provided, or when the caller requests + the first available PPISP camera attributes regardless of camera binding. Args: stage: USD stage to search. Returns: - Parsed :class:`PpispCfg` for the first matching shader, else ``None``. + Parsed :class:`PpispCfg` for the first matching camera, else ``None``. """ - from pxr import UsdShade - for prim in stage.Traverse(): - if prim.GetName() != PPISP_SHADER_NAME: - continue - shader = UsdShade.Shader(prim) - if shader: - return ppisp_cfg_from_usd_shader(shader) + if _has_ppisp_camera_attrs(prim): + cfg = ppisp_cfg_from_usd_camera(prim) + cfg.camera_prim_path = None + return cfg return None diff --git a/source/isaaclab_ppisp/isaaclab_ppisp/kernels.py b/source/isaaclab_ppisp/isaaclab_ppisp/kernels.py index f5ad5a29878b..02a3ed36cd7c 100644 --- a/source/isaaclab_ppisp/isaaclab_ppisp/kernels.py +++ b/source/isaaclab_ppisp/isaaclab_ppisp/kernels.py @@ -7,12 +7,59 @@ from __future__ import annotations +from string import Template + import warp as wp -from .cfg import PpispCfg +from .cfg import PPISP_CONTROLLER_EXPECTED_WEIGHTS_LEN, PpispCfg wp.init() +PPISP_CONTROLLER_INPUT_DOWNSAMPLING = 3 +PPISP_CONTROLLER_CNN_FEATURE_DIM = 64 +PPISP_CONTROLLER_CNN_FEATURE_CHUNK = 16 +PPISP_CONTROLLER_POOL_GRID_H = 5 +PPISP_CONTROLLER_POOL_GRID_W = 5 +PPISP_CONTROLLER_POOL_CELL_COUNT = PPISP_CONTROLLER_POOL_GRID_H * PPISP_CONTROLLER_POOL_GRID_W +PPISP_CONTROLLER_FEATURE_LEN = PPISP_CONTROLLER_CNN_FEATURE_DIM * PPISP_CONTROLLER_POOL_CELL_COUNT +PPISP_CONTROLLER_MLP_INPUT_DIM = PPISP_CONTROLLER_FEATURE_LEN + 1 +PPISP_CONTROLLER_HIDDEN_DIM = 128 +PPISP_CONTROLLER_PARAM_COUNT = 9 +PPISP_CONTROLLER_POOL_THREAD_GROUP_SIZE = 256 +PPISP_CONTROLLER_MLP_THREAD_GROUP_SIZE = 128 + +# This layout mirrors NRE's exported PPISP controller weight tensor layout. +# Linear/conv matrices are flattened as row-major ``[out_channel, in_channel]`` +# slices. +PPISP_CONTROLLER_OFF_CONV1_W = 0 +PPISP_CONTROLLER_OFF_CONV1_B = PPISP_CONTROLLER_OFF_CONV1_W + 16 * 3 +PPISP_CONTROLLER_OFF_CONV2_W = PPISP_CONTROLLER_OFF_CONV1_B + 16 +PPISP_CONTROLLER_OFF_CONV2_B = PPISP_CONTROLLER_OFF_CONV2_W + 32 * 16 +PPISP_CONTROLLER_OFF_CONV3_W = PPISP_CONTROLLER_OFF_CONV2_B + 32 +PPISP_CONTROLLER_OFF_CONV3_B = PPISP_CONTROLLER_OFF_CONV3_W + 64 * 32 +PPISP_CONTROLLER_OFF_TRUNK0_W = PPISP_CONTROLLER_OFF_CONV3_B + 64 +PPISP_CONTROLLER_OFF_TRUNK0_B = ( + PPISP_CONTROLLER_OFF_TRUNK0_W + PPISP_CONTROLLER_HIDDEN_DIM * PPISP_CONTROLLER_MLP_INPUT_DIM +) +PPISP_CONTROLLER_OFF_TRUNK1_W = PPISP_CONTROLLER_OFF_TRUNK0_B + PPISP_CONTROLLER_HIDDEN_DIM +PPISP_CONTROLLER_OFF_TRUNK1_B = ( + PPISP_CONTROLLER_OFF_TRUNK1_W + PPISP_CONTROLLER_HIDDEN_DIM * PPISP_CONTROLLER_HIDDEN_DIM +) +PPISP_CONTROLLER_OFF_TRUNK2_W = PPISP_CONTROLLER_OFF_TRUNK1_B + PPISP_CONTROLLER_HIDDEN_DIM +PPISP_CONTROLLER_OFF_TRUNK2_B = ( + PPISP_CONTROLLER_OFF_TRUNK2_W + PPISP_CONTROLLER_HIDDEN_DIM * PPISP_CONTROLLER_HIDDEN_DIM +) +PPISP_CONTROLLER_OFF_EXP_W = PPISP_CONTROLLER_OFF_TRUNK2_B + PPISP_CONTROLLER_HIDDEN_DIM +PPISP_CONTROLLER_OFF_EXP_B = PPISP_CONTROLLER_OFF_EXP_W + PPISP_CONTROLLER_HIDDEN_DIM +PPISP_CONTROLLER_OFF_COL_W = PPISP_CONTROLLER_OFF_EXP_B + 1 +PPISP_CONTROLLER_OFF_COL_B = PPISP_CONTROLLER_OFF_COL_W + 8 * PPISP_CONTROLLER_HIDDEN_DIM +PPISP_CONTROLLER_TOTAL_WEIGHTS = PPISP_CONTROLLER_OFF_COL_B + 8 +if PPISP_CONTROLLER_TOTAL_WEIGHTS != PPISP_CONTROLLER_EXPECTED_WEIGHTS_LEN: + raise RuntimeError( + "PPISP controller weight offsets do not match the expected exported weight count: " + f"{PPISP_CONTROLLER_TOTAL_WEIGHTS} != {PPISP_CONTROLLER_EXPECTED_WEIGHTS_LEN}." + ) + @wp.func def _bounded_softplus(raw: wp.float32, min_value: wp.float32): @@ -160,6 +207,307 @@ def _compute_homography_mul( return wp.vec3f(rgi[0], rgi[1], rgi[2] - rgi[0] - rgi[1]) +_PPISP_CONTROLLER_POOL_NATIVE_SNIPPET = Template(r""" + __shared__ float gs_reduce[$pool_thread_group_size]; + + const int input_downsampling = $input_downsampling; + const int cnn_feature_dim = $cnn_feature_dim; + const int cnn_feature_chunk = $cnn_feature_chunk; + const int pool_grid_h = $pool_grid_h; + const int pool_grid_w = $pool_grid_w; + const int pool_cell_count = $pool_cell_count; + const int pool_thread_group_size = $pool_thread_group_size; + const int off_conv1_w = $off_conv1_w; + const int off_conv1_b = $off_conv1_b; + const int off_conv2_w = $off_conv2_w; + const int off_conv2_b = $off_conv2_b; + const int off_conv3_w = $off_conv3_w; + const int off_conv3_b = $off_conv3_b; + + int ds_width = image_width / input_downsampling; + int ds_height = image_height / input_downsampling; + if (ds_width < 1) { + ds_width = 1; + } + if (ds_height < 1) { + ds_height = 1; + } + + const int gy = cell / pool_grid_w; + const int gx = cell - gy * pool_grid_w; + + const int h_start = (gy * ds_height) / pool_grid_h; + int h_end = ((gy + 1) * ds_height + pool_grid_h - 1) / pool_grid_h; + const int w_start = (gx * ds_width) / pool_grid_w; + int w_end = ((gx + 1) * ds_width + pool_grid_w - 1) / pool_grid_w; + if (h_end > ds_height) { + h_end = ds_height; + } + if (w_end > ds_width) { + w_end = ds_width; + } + + const int cell_width = (w_end > w_start) ? (w_end - w_start) : 0; + const int cell_height = (h_end > h_start) ? (h_end - h_start) : 0; + const int count = cell_width * cell_height; + + for (int first_channel = 0; first_channel < cnn_feature_dim; first_channel += cnn_feature_chunk) { + float partial[16]; +#pragma unroll + for (int c = 0; c < cnn_feature_chunk; ++c) { + partial[c] = 0.0f; + } + + for (int idx = lane; idx < count; idx += pool_thread_group_size) { + const int dy = h_start + idx / cell_width; + const int dx = w_start + idx - (idx / cell_width) * cell_width; + const int x0 = dx * input_downsampling; + const int y0 = dy * input_downsampling; + int x1 = x0 + input_downsampling; + int y1 = y0 + input_downsampling; + if (x1 > image_width) { + x1 = image_width; + } + if (y1 > image_height) { + y1 = image_height; + } + + float pooled[16]; +#pragma unroll + for (int c = 0; c < 16; ++c) { + pooled[c] = -3.4028234663852886e38f; + } + + for (int yy = y0; yy < y1; ++yy) { + for (int xx = x0; xx < x1; ++xx) { + const float r = (*wp::address(hdr_color, camera_id, yy, xx, 0)) * responsivity; + const float g = (*wp::address(hdr_color, camera_id, yy, xx, 1)) * responsivity; + const float b = (*wp::address(hdr_color, camera_id, yy, xx, 2)) * responsivity; + +#pragma unroll + for (int o = 0; o < 16; ++o) { + float v = *wp::address(weights, off_conv1_b + o); + v += r * (*wp::address(weights, off_conv1_w + o * 3 + 0)); + v += g * (*wp::address(weights, off_conv1_w + o * 3 + 1)); + v += b * (*wp::address(weights, off_conv1_w + o * 3 + 2)); + pooled[o] = fmaxf(pooled[o], v); + } + } + } + +#pragma unroll + for (int c = 0; c < 16; ++c) { + pooled[c] = fmaxf(0.0f, pooled[c]); + } + + float conv2[32]; +#pragma unroll + for (int o = 0; o < 32; ++o) { + float v = *wp::address(weights, off_conv2_b + o); +#pragma unroll + for (int i = 0; i < 16; ++i) { + v += pooled[i] * (*wp::address(weights, off_conv2_w + o * 16 + i)); + } + conv2[o] = fmaxf(0.0f, v); + } + +#pragma unroll + for (int c = 0; c < cnn_feature_chunk; ++c) { + const int o = first_channel + c; + float v = *wp::address(weights, off_conv3_b + o); +#pragma unroll + for (int i = 0; i < 32; ++i) { + v += conv2[i] * (*wp::address(weights, off_conv3_w + o * 32 + i)); + } + partial[c] += v; + } + } + +#pragma unroll + for (int c = 0; c < cnn_feature_chunk; ++c) { + gs_reduce[lane] = partial[c]; + __syncthreads(); + + for (int stride = pool_thread_group_size / 2; stride > 0; stride >>= 1) { + if (lane < stride) { + gs_reduce[lane] += gs_reduce[lane + stride]; + } + __syncthreads(); + } + + if (lane == 0) { + const int channel = first_channel + c; + // Channel-major layout matches the controller trunk0 feature ordering. + const int feature_id = channel * pool_cell_count + cell; + const float inv_count = (count > 0) ? (1.0f / ((float)count)) : 0.0f; + *wp::address(features, camera_id, feature_id) = gs_reduce[0] * inv_count; + } + __syncthreads(); + } + } +""").substitute( + { + "input_downsampling": PPISP_CONTROLLER_INPUT_DOWNSAMPLING, + "cnn_feature_dim": PPISP_CONTROLLER_CNN_FEATURE_DIM, + "cnn_feature_chunk": PPISP_CONTROLLER_CNN_FEATURE_CHUNK, + "pool_grid_h": PPISP_CONTROLLER_POOL_GRID_H, + "pool_grid_w": PPISP_CONTROLLER_POOL_GRID_W, + "pool_cell_count": PPISP_CONTROLLER_POOL_CELL_COUNT, + "pool_thread_group_size": PPISP_CONTROLLER_POOL_THREAD_GROUP_SIZE, + "off_conv1_w": PPISP_CONTROLLER_OFF_CONV1_W, + "off_conv1_b": PPISP_CONTROLLER_OFF_CONV1_B, + "off_conv2_w": PPISP_CONTROLLER_OFF_CONV2_W, + "off_conv2_b": PPISP_CONTROLLER_OFF_CONV2_B, + "off_conv3_w": PPISP_CONTROLLER_OFF_CONV3_W, + "off_conv3_b": PPISP_CONTROLLER_OFF_CONV3_B, + } +) + + +@wp.func_native(_PPISP_CONTROLLER_POOL_NATIVE_SNIPPET) +def _ppisp_controller_pool_features_native( + hdr_color: wp.array4d(dtype=wp.float32), + weights: wp.array(dtype=wp.float32), + features: wp.array2d(dtype=wp.float32), + image_width: wp.int32, + image_height: wp.int32, + responsivity: wp.float32, + camera_id: wp.int32, + cell: wp.int32, + lane: wp.int32, +): ... + + +@wp.kernel(enable_backward=False) +def _ppisp_controller_pool_features_native_kernel( + hdr_color: wp.array4d(dtype=wp.float32), + weights: wp.array(dtype=wp.float32), + features: wp.array2d(dtype=wp.float32), + image_width: wp.int32, + image_height: wp.int32, + responsivity: wp.float32, +): + """Run the controller CNN and adaptive pool using a native CUDA snippet.""" + camera_id, cell, lane = wp.tid() + _ppisp_controller_pool_features_native( + hdr_color, + weights, + features, + image_width, + image_height, + responsivity, + camera_id, + cell, + lane, + ) + + +_PPISP_CONTROLLER_MLP_NATIVE_SNIPPET = Template(r""" + __shared__ float hidden_a[$hidden_dim]; + __shared__ float hidden_b[$hidden_dim]; + + const int pool_feature_len = $feature_len; + const int mlp_input_dim = $mlp_input_dim; + const int mlp_hidden_dim = $hidden_dim; + const int color_params_per_frame = 8; + const int mlp_thread_group_size = $mlp_thread_group_size; + const int off_trunk0_w = $off_trunk0_w; + const int off_trunk0_b = $off_trunk0_b; + const int off_trunk1_w = $off_trunk1_w; + const int off_trunk1_b = $off_trunk1_b; + const int off_trunk2_w = $off_trunk2_w; + const int off_trunk2_b = $off_trunk2_b; + const int off_exp_w = $off_exp_w; + const int off_exp_b = $off_exp_b; + const int off_col_w = $off_col_w; + const int off_col_b = $off_col_b; + + for (int o = lane; o < mlp_hidden_dim; o += mlp_thread_group_size) { + float v = *wp::address(weights, off_trunk0_b + o); + for (int i = 0; i < pool_feature_len; ++i) { + v += (*wp::address(features, camera_id, i)) * (*wp::address(weights, off_trunk0_w + o * mlp_input_dim + i)); + } + v += prior_exposure * (*wp::address(weights, off_trunk0_w + o * mlp_input_dim + pool_feature_len)); + hidden_a[o] = fmaxf(0.0f, v); + } + __syncthreads(); + + for (int o = lane; o < mlp_hidden_dim; o += mlp_thread_group_size) { + float v = *wp::address(weights, off_trunk1_b + o); + for (int i = 0; i < mlp_hidden_dim; ++i) { + v += hidden_a[i] * (*wp::address(weights, off_trunk1_w + o * mlp_hidden_dim + i)); + } + hidden_b[o] = fmaxf(0.0f, v); + } + __syncthreads(); + + for (int o = lane; o < mlp_hidden_dim; o += mlp_thread_group_size) { + float v = *wp::address(weights, off_trunk2_b + o); + for (int i = 0; i < mlp_hidden_dim; ++i) { + v += hidden_b[i] * (*wp::address(weights, off_trunk2_w + o * mlp_hidden_dim + i)); + } + hidden_a[o] = fmaxf(0.0f, v); + } + __syncthreads(); + + if (lane == 0) { + float v = *wp::address(weights, off_exp_b); + for (int i = 0; i < mlp_hidden_dim; ++i) { + v += hidden_a[i] * (*wp::address(weights, off_exp_w + i)); + } + *wp::address(controller_params, camera_id, 0) = v; + } + if (lane < color_params_per_frame) { + const int o = lane; + float v = *wp::address(weights, off_col_b + o); + for (int i = 0; i < mlp_hidden_dim; ++i) { + v += hidden_a[i] * (*wp::address(weights, off_col_w + o * mlp_hidden_dim + i)); + } + *wp::address(controller_params, camera_id, 1 + o) = v; + } +""").substitute( + { + "feature_len": PPISP_CONTROLLER_FEATURE_LEN, + "mlp_input_dim": PPISP_CONTROLLER_MLP_INPUT_DIM, + "hidden_dim": PPISP_CONTROLLER_HIDDEN_DIM, + "mlp_thread_group_size": PPISP_CONTROLLER_MLP_THREAD_GROUP_SIZE, + "off_trunk0_w": PPISP_CONTROLLER_OFF_TRUNK0_W, + "off_trunk0_b": PPISP_CONTROLLER_OFF_TRUNK0_B, + "off_trunk1_w": PPISP_CONTROLLER_OFF_TRUNK1_W, + "off_trunk1_b": PPISP_CONTROLLER_OFF_TRUNK1_B, + "off_trunk2_w": PPISP_CONTROLLER_OFF_TRUNK2_W, + "off_trunk2_b": PPISP_CONTROLLER_OFF_TRUNK2_B, + "off_exp_w": PPISP_CONTROLLER_OFF_EXP_W, + "off_exp_b": PPISP_CONTROLLER_OFF_EXP_B, + "off_col_w": PPISP_CONTROLLER_OFF_COL_W, + "off_col_b": PPISP_CONTROLLER_OFF_COL_B, + } +) + + +@wp.func_native(_PPISP_CONTROLLER_MLP_NATIVE_SNIPPET) +def _ppisp_controller_mlp_native( + features: wp.array2d(dtype=wp.float32), + weights: wp.array(dtype=wp.float32), + controller_params: wp.array2d(dtype=wp.float32), + prior_exposure: wp.float32, + camera_id: wp.int32, + lane: wp.int32, +): ... + + +@wp.kernel(enable_backward=False) +def _ppisp_controller_mlp_native_kernel( + features: wp.array2d(dtype=wp.float32), + weights: wp.array(dtype=wp.float32), + controller_params: wp.array2d(dtype=wp.float32), + prior_exposure: wp.float32, +): + """Run the controller MLP using a native CUDA snippet.""" + camera_id, lane = wp.tid() + _ppisp_controller_mlp_native(features, weights, controller_params, prior_exposure, camera_id, lane) + + @wp.kernel(enable_backward=False) def _apply_ppisp_kernel( hdr_color: wp.array4d(dtype=wp.float32), @@ -246,6 +594,85 @@ def _apply_ppisp_kernel( out_rgba[camera_id, height_id, width_id, 3] = wp.uint8(255) +@wp.kernel(enable_backward=False) +def _apply_ppisp_controller_kernel( + hdr_color: wp.array4d(dtype=wp.float32), + out_rgba: wp.array4d(dtype=wp.uint8), + controller_params: wp.array2d(dtype=wp.float32), + image_width: wp.int32, + image_height: wp.int32, + responsivity: wp.float32, + vignetting_center_r: wp.vec2f, + vignetting_alpha1_r: wp.float32, + vignetting_alpha2_r: wp.float32, + vignetting_alpha3_r: wp.float32, + vignetting_center_g: wp.vec2f, + vignetting_alpha1_g: wp.float32, + vignetting_alpha2_g: wp.float32, + vignetting_alpha3_g: wp.float32, + vignetting_center_b: wp.vec2f, + vignetting_alpha1_b: wp.float32, + vignetting_alpha2_b: wp.float32, + vignetting_alpha3_b: wp.float32, + crf_toe_r: wp.float32, + crf_shoulder_r: wp.float32, + crf_gamma_r: wp.float32, + crf_center_r: wp.float32, + crf_toe_g: wp.float32, + crf_shoulder_g: wp.float32, + crf_gamma_g: wp.float32, + crf_center_g: wp.float32, + crf_toe_b: wp.float32, + crf_shoulder_b: wp.float32, + crf_gamma_b: wp.float32, + crf_center_b: wp.float32, +): + """Apply PPISP using per-camera exposure/color parameters predicted by the controller.""" + camera_id, height_id, width_id = wp.tid() + rgb = wp.vec3f( + hdr_color[camera_id, height_id, width_id, 0], + hdr_color[camera_id, height_id, width_id, 1], + hdr_color[camera_id, height_id, width_id, 2], + ) + max_resolution = wp.float32(image_width) + if image_height > image_width: + max_resolution = wp.float32(image_height) + uv = wp.vec2f( + (wp.float32(width_id) + 0.5 - wp.float32(image_width) * 0.5) / max_resolution, + (wp.float32(height_id) + 0.5 - wp.float32(image_height) * 0.5) / max_resolution, + ) + + # Exported controller params are: + # [exposure, blue.x, blue.y, red.x, red.y, green.x, green.y, neutral.x, neutral.y]. + color_latent_blue = wp.vec2f(controller_params[camera_id, 1], controller_params[camera_id, 2]) + color_latent_red = wp.vec2f(controller_params[camera_id, 3], controller_params[camera_id, 4]) + color_latent_green = wp.vec2f(controller_params[camera_id, 5], controller_params[camera_id, 6]) + color_latent_neutral = wp.vec2f(controller_params[camera_id, 7], controller_params[camera_id, 8]) + + out_rgb = rgb * responsivity + out_rgb = out_rgb * wp.pow(2.0, controller_params[camera_id, 0]) + out_rgb[0] = _apply_vignetting( + out_rgb[0], uv, vignetting_center_r, vignetting_alpha1_r, vignetting_alpha2_r, vignetting_alpha3_r + ) + out_rgb[1] = _apply_vignetting( + out_rgb[1], uv, vignetting_center_g, vignetting_alpha1_g, vignetting_alpha2_g, vignetting_alpha3_g + ) + out_rgb[2] = _apply_vignetting( + out_rgb[2], uv, vignetting_center_b, vignetting_alpha1_b, vignetting_alpha2_b, vignetting_alpha3_b + ) + out_rgb = _compute_homography_mul( + out_rgb, color_latent_blue, color_latent_red, color_latent_green, color_latent_neutral + ) + out_rgb[0] = _apply_crf(out_rgb[0], crf_toe_r, crf_shoulder_r, crf_gamma_r, crf_center_r) + out_rgb[1] = _apply_crf(out_rgb[1], crf_toe_g, crf_shoulder_g, crf_gamma_g, crf_center_g) + out_rgb[2] = _apply_crf(out_rgb[2], crf_toe_b, crf_shoulder_b, crf_gamma_b, crf_center_b) + + out_rgba[camera_id, height_id, width_id, 0] = wp.uint8(wp.clamp(out_rgb[0], 0.0, 1.0) * 255.0) + out_rgba[camera_id, height_id, width_id, 1] = wp.uint8(wp.clamp(out_rgb[1], 0.0, 1.0) * 255.0) + out_rgba[camera_id, height_id, width_id, 2] = wp.uint8(wp.clamp(out_rgb[2], 0.0, 1.0) * 255.0) + out_rgba[camera_id, height_id, width_id, 3] = wp.uint8(255) + + def apply_ppisp_to_rgba(hdr_color: wp.array, out_rgba: wp.array, cfg: PpispCfg) -> None: """Apply PPISP to ``hdr_color`` and write LDR RGBA into ``out_rgba``. @@ -303,3 +730,141 @@ def apply_ppisp_to_rgba(hdr_color: wp.array, out_rgba: wp.array, cfg: PpispCfg) ], device=str(out_rgba.device), ) + + +def _validate_ppisp_controller_inputs( + hdr_color: wp.array, + controller_weights: wp.array, + features: wp.array, + controller_params: wp.array, +) -> None: + if hdr_color.dtype is not wp.float32: + raise ValueError(f"Camera PPISP controller HDR input must be wp.float32, got {hdr_color.dtype}.") + if controller_weights.dtype is not wp.float32: + raise ValueError(f"Camera PPISP controller weights must be wp.float32, got {controller_weights.dtype}.") + if controller_weights.shape[0] != PPISP_CONTROLLER_EXPECTED_WEIGHTS_LEN: + raise ValueError( + "Camera PPISP controller weights must have shape " + f"({PPISP_CONTROLLER_EXPECTED_WEIGHTS_LEN},), got {controller_weights.shape}." + ) + if features.shape != (hdr_color.shape[0], PPISP_CONTROLLER_FEATURE_LEN): + raise ValueError( + "Camera PPISP controller features must have shape " + f"({hdr_color.shape[0]}, {PPISP_CONTROLLER_FEATURE_LEN}), got {features.shape}." + ) + if controller_params.shape != (hdr_color.shape[0], PPISP_CONTROLLER_PARAM_COUNT): + raise ValueError( + "Camera PPISP controller params must have shape " + f"({hdr_color.shape[0]}, {PPISP_CONTROLLER_PARAM_COUNT}), got {controller_params.shape}." + ) + + +def compute_ppisp_controller_params( + hdr_color: wp.array, + controller_weights: wp.array, + features: wp.array, + controller_params: wp.array, + prior_exposure: float, + responsivity: float = 1.0, +) -> None: + """Run the PPISP controller using native CUDA snippets inside Warp kernels. + + Args: + hdr_color: HDR scene-linear input, shape ``(N, H, W, 3)``. + controller_weights: Flattened embedded controller weights. + features: Scratch buffer, shape ``(N, 1600)``. + controller_params: Output buffer, shape ``(N, 9)``. + prior_exposure: Scalar prior exposure supplied by the PPISP config. + responsivity: Achromatic HDR multiplier applied before controller + feature extraction. + + Raises: + ValueError: When ``hdr_color`` is not on a CUDA device. + """ + _validate_ppisp_controller_inputs(hdr_color, controller_weights, features, controller_params) + device = wp.get_device(str(hdr_color.device)) + if not device.is_cuda: + raise ValueError("Camera PPISP controller requires a CUDA device.") + + image_height = int(hdr_color.shape[1]) + image_width = int(hdr_color.shape[2]) + device_name = str(hdr_color.device) + wp.launch_tiled( + _ppisp_controller_pool_features_native_kernel, + dim=(int(hdr_color.shape[0]), PPISP_CONTROLLER_POOL_CELL_COUNT), + inputs=[hdr_color, controller_weights, features, image_width, image_height, float(responsivity)], + device=device_name, + block_dim=PPISP_CONTROLLER_POOL_THREAD_GROUP_SIZE, + ) + wp.launch_tiled( + _ppisp_controller_mlp_native_kernel, + dim=(int(hdr_color.shape[0]),), + inputs=[features, controller_weights, controller_params, float(prior_exposure)], + device=device_name, + block_dim=PPISP_CONTROLLER_MLP_THREAD_GROUP_SIZE, + ) + + +def apply_ppisp_to_rgba_with_controller_params( + hdr_color: wp.array, out_rgba: wp.array, cfg: PpispCfg, controller_params: wp.array +) -> None: + """Apply PPISP with per-camera controller parameters. + + Args: + hdr_color: HDR scene-linear input, shape ``(N, H, W, 3)``. + out_rgba: LDR RGBA output, shape ``(N, H, W, 4)``. + cfg: PPISP configuration. Static inputs provide responsivity, + vignetting and CRF. + controller_params: ``wp.array2d(float32)`` of shape ``(N, 9)`` holding + ``[exposureOffset, colorLatentBlue.xy, colorLatentRed.xy, + colorLatentGreen.xy, colorLatentNeutral.xy]``. + """ + if hdr_color.dtype is not wp.float32: + raise ValueError(f"Camera PPISP HDR input must be wp.float32, got {hdr_color.dtype}.") + if out_rgba.dtype is not wp.uint8: + raise ValueError(f"Camera PPISP RGBA output must be wp.uint8, got {out_rgba.dtype}.") + if controller_params.dtype is not wp.float32: + raise ValueError(f"Camera PPISP controller params must be wp.float32, got {controller_params.dtype}.") + if controller_params.shape[0] != out_rgba.shape[0] or controller_params.shape[1] != 9: + raise ValueError( + f"Camera PPISP controller params must have shape ({out_rgba.shape[0]}, 9), got {controller_params.shape}." + ) + + inputs = cfg.inputs + wp.launch( + _apply_ppisp_controller_kernel, + dim=out_rgba.shape[:3], + inputs=[ + hdr_color, + out_rgba, + controller_params, + int(out_rgba.shape[2]), + int(out_rgba.shape[1]), + float(inputs.get("responsivity", 1.0)), + wp.vec2f(*inputs["vignettingCenterR"]), + float(inputs["vignettingAlpha1R"]), + float(inputs["vignettingAlpha2R"]), + float(inputs["vignettingAlpha3R"]), + wp.vec2f(*inputs["vignettingCenterG"]), + float(inputs["vignettingAlpha1G"]), + float(inputs["vignettingAlpha2G"]), + float(inputs["vignettingAlpha3G"]), + wp.vec2f(*inputs["vignettingCenterB"]), + float(inputs["vignettingAlpha1B"]), + float(inputs["vignettingAlpha2B"]), + float(inputs["vignettingAlpha3B"]), + float(inputs["crfToeR"]), + float(inputs["crfShoulderR"]), + float(inputs["crfGammaR"]), + float(inputs["crfCenterR"]), + float(inputs["crfToeG"]), + float(inputs["crfShoulderG"]), + float(inputs["crfGammaG"]), + float(inputs["crfCenterG"]), + float(inputs["crfToeB"]), + float(inputs["crfShoulderB"]), + float(inputs["crfGammaB"]), + float(inputs["crfCenterB"]), + ], + device=str(out_rgba.device), + ) diff --git a/source/isaaclab_ppisp/isaaclab_ppisp/pipeline.py b/source/isaaclab_ppisp/isaaclab_ppisp/pipeline.py index f1593df176ae..345215c8c9b2 100644 --- a/source/isaaclab_ppisp/isaaclab_ppisp/pipeline.py +++ b/source/isaaclab_ppisp/isaaclab_ppisp/pipeline.py @@ -14,12 +14,16 @@ from __future__ import annotations -from typing import Any - import warp as wp from .cfg import PpispCfg, normalize_ppisp_cfg -from .kernels import apply_ppisp_to_rgba +from .kernels import ( + PPISP_CONTROLLER_FEATURE_LEN, + PPISP_CONTROLLER_PARAM_COUNT, + apply_ppisp_to_rgba, + apply_ppisp_to_rgba_with_controller_params, + compute_ppisp_controller_params, +) class PpispPipeline: @@ -29,33 +33,83 @@ class PpispPipeline: normalised :class:`PpispCfg` and dispatches the PPISP Warp kernel once per render tick via :meth:`apply`. - One pipeline instance applies to the whole Camera sensor batch. The PPISP - Warp kernel takes scalar coefficients, so every cloned view in a tiled - batch shares the same ISP configuration — there is no per-view ISP today. - Per-view support would require packing the cfg into GPU arrays and indexing - by ``camera_id`` inside the kernel. + One pipeline instance applies to the whole Camera sensor batch. Static + PPISP coefficients are scalar values shared by every cloned view in a + tiled batch. When controller weights are configured, the controller pass + instead predicts exposure and color-latent values independently for each + camera view from its HDR image; the image pass then combines those + per-view values with the shared static coefficients. Today only :class:`PpispCfg` is accepted; future ISP implementations can either subclass or be selected by cfg type without changes to the backend renderers. """ - def __init__(self, cfg: PpispCfg, stage: Any = None): + def __init__(self, cfg: PpispCfg): """Initialize the PPISP pipeline. Normalises ``cfg`` on construction (validates input keys, fills - defaults, and — when ``cfg.shader_prim_path`` is set and ``stage`` is - non-``None`` — merges shader-authored values with user overrides). + defaults). :class:`~isaaclab.sensors.camera.Camera` already normalises ``isp_cfg`` before passing the :class:`~isaaclab.renderers.CameraRenderSpec` to - the backend, so renderer backends typically pass ``stage=None`` here. + the backend, so renderer backends pass a concrete, resolved config here. Args: cfg: The PPISP configuration. - stage: Optional USD stage used to resolve ``cfg.shader_prim_path``. """ - self.cfg = normalize_ppisp_cfg(cfg, stage=stage) + normalized_cfg = normalize_ppisp_cfg(cfg) + if normalized_cfg is None: + raise ValueError("PpispPipeline requires a concrete PpispCfg.") + self.cfg = normalized_cfg + self._controller_weights_by_device: dict[str, wp.array] = {} + self._controller_buffers_by_shape: dict[tuple[str, int, int, int], tuple[wp.array, ...]] = {} def apply(self, hdr: wp.array, rgba: wp.array) -> None: """Run the PPISP kernel: HDR scene-linear → LDR RGBA, in place on ``rgba``.""" - apply_ppisp_to_rgba(hdr, rgba, self.cfg) + if self.cfg.controller_weights is None: + apply_ppisp_to_rgba(hdr, rgba, self.cfg) + return + controller_params = self._compute_controller_params(hdr) + apply_ppisp_to_rgba_with_controller_params(hdr, rgba, self.cfg, controller_params) + + def _compute_controller_params(self, hdr: wp.array) -> wp.array: + """Run the exported PPISP controller and return a Warp view of ``(N, 9)`` params.""" + controller_weights = self.cfg.controller_weights + assert controller_weights is not None + + device = str(hdr.device) + weights = self._controller_weights_by_device.get(device) + if weights is None: + weights = wp.array(controller_weights, dtype=wp.float32, device=device) + self._controller_weights_by_device[device] = weights + + features, controller_params = self._controller_buffers(hdr) + compute_ppisp_controller_params( + hdr, + weights, + features, + controller_params, + self.cfg.controller_prior_exposure, + float(self.cfg.controller_responsivity), + ) + return controller_params + + def _controller_buffers(self, hdr: wp.array) -> tuple[wp.array, ...]: + """Return cached controller scratch buffers matching ``hdr`` shape/device.""" + num_cameras = int(hdr.shape[0]) + image_height = int(hdr.shape[1]) + image_width = int(hdr.shape[2]) + device = str(hdr.device) + key = (device, num_cameras, image_height, image_width) + buffers = self._controller_buffers_by_shape.get(key) + if buffers is not None: + return buffers + + # Camera output shape is fixed for a pipeline lifetime; cache buffers so + # controller execution does not allocate every frame. + buffers = ( + wp.empty((num_cameras, PPISP_CONTROLLER_FEATURE_LEN), dtype=wp.float32, device=device), + wp.empty((num_cameras, PPISP_CONTROLLER_PARAM_COUNT), dtype=wp.float32, device=device), + ) + self._controller_buffers_by_shape[key] = buffers + return buffers diff --git a/source/isaaclab_ppisp/test/test_ppisp.py b/source/isaaclab_ppisp/test/test_ppisp.py index f9ca5ed19256..5b65ef570725 100644 --- a/source/isaaclab_ppisp/test/test_ppisp.py +++ b/source/isaaclab_ppisp/test/test_ppisp.py @@ -5,53 +5,126 @@ """Tests for PPISP USD parsing helpers.""" +from __future__ import annotations + import pytest -from isaaclab_ppisp import PpispCfg, auto_camera_ppisp_cfg, normalize_ppisp_cfg, ppisp_cfg_from_usd_shader +from isaaclab_ppisp import ( + PpispCfg, + auto_any_ppisp_cfg, + auto_camera_ppisp_cfg, + default_ppisp_inputs, + has_ppisp_camera_attrs, + normalize_ppisp_cfg, + ppisp_cfg_from_usd_camera, +) +from isaaclab_ppisp.cfg import PPISP_CONTROLLER_EXPECTED_WEIGHTS_LEN, resolve_and_normalize + +from pxr import Gf, Sdf, Usd, Vt + +_PPISP_FLOAT2_ATTRS = { + "vignettingCenterR", + "vignettingCenterG", + "vignettingCenterB", + "colorLatentBlue", + "colorLatentRed", + "colorLatentGreen", + "colorLatentNeutral", +} + + +def _author_ppisp_attr(camera_prim: Usd.Prim, name: str, value): + value_type = Sdf.ValueTypeNames.Float2 if name in _PPISP_FLOAT2_ATTRS else Sdf.ValueTypeNames.Float + attr = camera_prim.CreateAttribute(f"ppisp:{name}", value_type) + attr.Set(Gf.Vec2f(*value) if name in _PPISP_FLOAT2_ATTRS else value) + return attr + + +def _author_camera(stage: Usd.Stage, camera_path: str = "/World/Camera") -> Usd.Prim: + return stage.DefinePrim(camera_path, "Camera") + + +def _author_ppisp_camera( + stage: Usd.Stage, + camera_path: str = "/World/Camera_ppisp", + *, + inherits: str | None = "/World/Camera", + attrs: dict | None = None, + controller_weights: list[float] | None = None, +) -> Usd.Prim: + camera_prim = _author_camera(stage, camera_path) + if inherits is not None: + _author_camera(stage, inherits) + camera_prim.GetInherits().AddInherit(Sdf.Path(inherits)) -from pxr import Gf, Sdf, Usd, UsdShade + for name, value in (attrs or {}).items(): + _author_ppisp_attr(camera_prim, name, value) + if controller_weights is not None: + camera_prim.CreateAttribute("ppisp:controllerWeights", Sdf.ValueTypeNames.FloatArray).Set( + Vt.FloatArray(controller_weights) + ) + return camera_prim -def test_ppisp_shader_import_uses_first_time_sample(): + +def _controller_weights() -> list[float]: + return [0.0] * PPISP_CONTROLLER_EXPECTED_WEIGHTS_LEN + + +def test_ppisp_camera_attr_import_uses_first_time_sample(): stage = Usd.Stage.CreateInMemory() - shader = UsdShade.Shader.Define(stage, "/Render/RenderProduct/PPISP") + ppisp_camera = _author_ppisp_camera(stage, attrs={}) - exposure = shader.CreateInput("exposureOffset", Sdf.ValueTypeNames.Float).GetAttr() + exposure = ppisp_camera.CreateAttribute("ppisp:exposureOffset", Sdf.ValueTypeNames.Float) exposure.Set(1.0) exposure.Set(2.0, 10.0) exposure.Set(3.0, 20.0) - color = shader.CreateInput("colorLatentBlue", Sdf.ValueTypeNames.Float2).GetAttr() + color = ppisp_camera.CreateAttribute("ppisp:colorLatentBlue", Sdf.ValueTypeNames.Float2) color.Set(Gf.Vec2f(0.0, 0.0)) color.Set(Gf.Vec2f(0.1, 0.2), 5.0) - cfg = ppisp_cfg_from_usd_shader(shader) + cfg = ppisp_cfg_from_usd_camera(ppisp_camera) + assert cfg.camera_prim_path is None assert cfg.inputs["exposureOffset"] == 2.0 assert cfg.inputs["colorLatentBlue"] == pytest.approx((0.1, 0.2)) -def test_normalize_ppisp_cfg_imports_shader_prim_path_from_stage(): +def test_normalize_ppisp_cfg_imports_camera_attrs_from_stage(): stage = Usd.Stage.CreateInMemory() - shader = UsdShade.Shader.Define(stage, "/Render/RenderProduct/PPISP") - shader.CreateInput("exposureOffset", Sdf.ValueTypeNames.Float).Set(1.5) - shader.CreateInput("colorLatentRed", Sdf.ValueTypeNames.Float2).Set(Gf.Vec2f(0.25, -0.5)) + _author_ppisp_camera( + stage, + attrs={ + "exposureOffset": 1.5, + "colorLatentRed": (0.25, -0.5), + }, + ) - cfg = normalize_ppisp_cfg(PpispCfg(shader_prim_path="/Render/RenderProduct/PPISP"), stage=stage) + cfg = normalize_ppisp_cfg(PpispCfg(camera_prim_path="/World/Camera_ppisp"), stage=stage) - assert cfg.shader_prim_path == "/Render/RenderProduct/PPISP" + assert cfg.camera_prim_path is None assert cfg.inputs["exposureOffset"] == 1.5 assert cfg.inputs["colorLatentRed"] == pytest.approx((0.25, -0.5)) -def test_normalize_ppisp_cfg_applies_explicit_overrides_after_shader_import(): +def test_normalize_ppisp_cfg_requires_stage_for_camera_prim_path(): + with pytest.raises(ValueError, match="requires a USD stage"): + normalize_ppisp_cfg(PpispCfg(camera_prim_path="/World/Camera_ppisp")) + + +def test_normalize_ppisp_cfg_applies_explicit_overrides_after_camera_attr_import(): stage = Usd.Stage.CreateInMemory() - shader = UsdShade.Shader.Define(stage, "/Render/RenderProduct/PPISP") - shader.CreateInput("exposureOffset", Sdf.ValueTypeNames.Float).Set(1.5) - shader.CreateInput("colorLatentRed", Sdf.ValueTypeNames.Float2).Set(Gf.Vec2f(0.25, -0.5)) + _author_ppisp_camera( + stage, + attrs={ + "exposureOffset": 1.5, + "colorLatentRed": (0.25, -0.5), + }, + ) cfg = normalize_ppisp_cfg( PpispCfg( - shader_prim_path="/Render/RenderProduct/PPISP", + camera_prim_path="/World/Camera_ppisp", inputs={"exposureOffset": 2.0}, ), stage=stage, @@ -61,28 +134,157 @@ def test_normalize_ppisp_cfg_applies_explicit_overrides_after_shader_import(): assert cfg.inputs["colorLatentRed"] == pytest.approx((0.25, -0.5)) -def test_auto_camera_ppisp_cfg_skips_matching_render_product_without_ppisp(): - """Camera-bound discovery skips generated RenderProducts that lack PPISP. +def test_normalize_ppisp_cfg_default_list_values_do_not_override_camera_attrs(): + stage = Usd.Stage.CreateInMemory() + _author_ppisp_camera(stage, attrs={"exposureOffset": 1.5}) + inputs = default_ppisp_inputs() + inputs["vignettingCenterR"] = [0.0, 0.0] + + cfg = normalize_ppisp_cfg(PpispCfg(camera_prim_path="/World/Camera_ppisp", inputs=inputs), stage=stage) + + assert cfg.inputs["exposureOffset"] == pytest.approx(1.5) + + +def test_normalize_ppisp_cfg_controller_responsivity_follows_responsivity_override(): + stage = Usd.Stage.CreateInMemory() + _author_ppisp_camera( + stage, + attrs={"responsivity": 2.5}, + controller_weights=_controller_weights(), + ) + + cfg = normalize_ppisp_cfg( + PpispCfg( + camera_prim_path="/World/Camera_ppisp", + inputs={"responsivity": 3.0}, + ), + stage=stage, + ) + + assert cfg.inputs["responsivity"] == pytest.approx(3.0) + assert cfg.controller_responsivity == pytest.approx(3.0) + + +def test_ppisp_cfg_from_usd_camera_requires_camera_attrs(): + stage = Usd.Stage.CreateInMemory() + camera = _author_camera(stage) + + with pytest.raises(ValueError, match="expected ppisp:\\* attributes"): + ppisp_cfg_from_usd_camera(camera) + + +def test_has_ppisp_camera_attrs_ignores_unknown_ppisp_attrs(): + stage = Usd.Stage.CreateInMemory() + camera = _author_camera(stage) + camera.CreateAttribute("ppisp:version", Sdf.ValueTypeNames.String).Set("1") + + assert not has_ppisp_camera_attrs(camera) + + +def test_auto_camera_ppisp_cfg_reads_direct_camera_attrs(): + stage = Usd.Stage.CreateInMemory() + _author_ppisp_camera( + stage, + "/World/Camera", + inherits=None, + attrs={ + "responsivity": 0.75, + "exposureOffset": 1.25, + "colorLatentBlue": (0.1, 0.2), + }, + ) + + cfg = auto_camera_ppisp_cfg(stage, "/World/Camera") + + assert cfg is not None + assert cfg.camera_prim_path is None + assert cfg.inputs["responsivity"] == pytest.approx(0.75) + assert cfg.inputs["exposureOffset"] == pytest.approx(1.25) + assert cfg.inputs["colorLatentBlue"] == pytest.approx((0.1, 0.2)) + - Isaac RTX can create a transient RenderProduct that targets the same camera - as an authored scene RenderProduct. If the transient prim has no ``PPISP`` - child, discovery must keep scanning instead of concluding the camera has no - camera-bound PPISP shader. - """ +def test_auto_camera_ppisp_cfg_does_not_scan_unmatched_cameras(): stage = Usd.Stage.CreateInMemory() - camera = "/World/Camera" - stage.DefinePrim(camera, "Camera") + _author_camera(stage, "/World/Camera") + _author_ppisp_camera( + stage, + "/World/Camera_ppisp", + inherits=None, + attrs={"exposureOffset": 1.5}, + ) - plain_rp = stage.DefinePrim("/Render/GeneratedRenderProduct", "RenderProduct") - plain_rp.CreateRelationship("camera").SetTargets([Sdf.Path(camera)]) + cfg = auto_camera_ppisp_cfg(stage, "/World/Camera") - authored_rp = stage.DefinePrim("/World/Render/RenderProduct", "RenderProduct") - authored_rp.CreateRelationship("camera").SetTargets([Sdf.Path(camera)]) - shader = UsdShade.Shader.Define(stage, "/World/Render/RenderProduct/PPISP") - shader.CreateInput("exposureOffset", Sdf.ValueTypeNames.Float).Set(1.5) + assert cfg is None - cfg = auto_camera_ppisp_cfg(stage, camera) + +def test_auto_any_ppisp_cfg_reads_first_camera_with_ppisp_attrs(): + stage = Usd.Stage.CreateInMemory() + _author_camera(stage, "/World/CameraWithoutPpisp") + _author_ppisp_camera( + stage, + "/World/CameraB_ppisp", + inherits=None, + attrs={"exposureOffset": 2.0}, + ) + _author_ppisp_camera( + stage, + "/World/CameraC_ppisp", + inherits=None, + attrs={"exposureOffset": 3.0}, + ) + + cfg = auto_any_ppisp_cfg(stage) assert cfg is not None - assert cfg.shader_prim_path == "/World/Render/RenderProduct/PPISP" - assert cfg.inputs["exposureOffset"] == 1.5 + assert cfg.camera_prim_path is None + assert cfg.inputs["exposureOffset"] == pytest.approx(2.0) + + +def test_resolve_and_normalize_without_camera_uses_first_ppisp_camera(): + from isaaclab.sensors.camera.camera_isp import CameraISPMode + + stage = Usd.Stage.CreateInMemory() + _author_camera(stage, "/World/CameraWithoutPpisp") + _author_ppisp_camera(stage, "/World/Camera_ppisp", inherits=None, attrs={"exposureOffset": 2.0}) + + cfg = resolve_and_normalize(CameraISPMode.AUTO_CAMERA, stage) + + assert cfg is not None + assert cfg.camera_prim_path is None + assert cfg.inputs["exposureOffset"] == pytest.approx(2.0) + + +def test_ppisp_cfg_from_usd_camera_reads_controller_weights_from_camera_attrs(): + stage = Usd.Stage.CreateInMemory() + camera = _author_ppisp_camera( + stage, + attrs={ + "responsivity": 2.5, + "vignettingAlpha1R": 0.25, + "crfToeB": 0.125, + }, + controller_weights=_controller_weights(), + ) + + cfg = ppisp_cfg_from_usd_camera(camera) + + assert cfg.inputs["responsivity"] == pytest.approx(2.5) + assert cfg.inputs["vignettingAlpha1R"] == pytest.approx(0.25) + assert cfg.inputs["crfToeB"] == pytest.approx(0.125) + assert cfg.controller_prior_exposure == pytest.approx(0.0) + assert cfg.controller_responsivity == pytest.approx(2.5) + assert cfg.controller_weights is not None + assert len(cfg.controller_weights) == PPISP_CONTROLLER_EXPECTED_WEIGHTS_LEN + + +def test_ppisp_cfg_from_usd_camera_validates_controller_weights_len(): + stage = Usd.Stage.CreateInMemory() + camera = _author_ppisp_camera( + stage, + attrs={"responsivity": 2.5}, + controller_weights=[0.0], + ) + + with pytest.raises(ValueError, match=f"Expected {PPISP_CONTROLLER_EXPECTED_WEIGHTS_LEN}"): + ppisp_cfg_from_usd_camera(camera) diff --git a/source/isaaclab_ppisp/test/test_ppisp_kernels.py b/source/isaaclab_ppisp/test/test_ppisp_kernels.py index 18ce05b0fc05..90f10a2ec3d1 100644 --- a/source/isaaclab_ppisp/test/test_ppisp_kernels.py +++ b/source/isaaclab_ppisp/test/test_ppisp_kernels.py @@ -3,8 +3,39 @@ # # SPDX-License-Identifier: BSD-3-Clause +import pytest import warp as wp -from isaaclab_ppisp import PpispCfg, apply_ppisp_to_rgba, normalize_ppisp_cfg +from isaaclab_ppisp import ( + PpispCfg, + apply_ppisp_to_rgba, + apply_ppisp_to_rgba_with_controller_params, + compute_ppisp_controller_params, + normalize_ppisp_cfg, +) +from isaaclab_ppisp.kernels import ( + PPISP_CONTROLLER_EXPECTED_WEIGHTS_LEN, + PPISP_CONTROLLER_FEATURE_LEN, + PPISP_CONTROLLER_HIDDEN_DIM, + PPISP_CONTROLLER_INPUT_DOWNSAMPLING, + PPISP_CONTROLLER_MLP_INPUT_DIM, + PPISP_CONTROLLER_OFF_COL_B, + PPISP_CONTROLLER_OFF_COL_W, + PPISP_CONTROLLER_OFF_CONV1_B, + PPISP_CONTROLLER_OFF_CONV1_W, + PPISP_CONTROLLER_OFF_CONV2_B, + PPISP_CONTROLLER_OFF_CONV2_W, + PPISP_CONTROLLER_OFF_CONV3_B, + PPISP_CONTROLLER_OFF_CONV3_W, + PPISP_CONTROLLER_OFF_EXP_B, + PPISP_CONTROLLER_OFF_EXP_W, + PPISP_CONTROLLER_OFF_TRUNK0_B, + PPISP_CONTROLLER_OFF_TRUNK0_W, + PPISP_CONTROLLER_OFF_TRUNK1_B, + PPISP_CONTROLLER_OFF_TRUNK1_W, + PPISP_CONTROLLER_OFF_TRUNK2_B, + PPISP_CONTROLLER_OFF_TRUNK2_W, + PPISP_CONTROLLER_PARAM_COUNT, +) from isaaclab.sensors.camera.tiled_camera_cfg import TiledCameraCfg @@ -19,6 +50,75 @@ def _rgba(shape: tuple[int, ...]) -> wp.array: return wp.zeros(shape=shape, dtype=wp.uint8) +def _controller_reference_params(hdr, weights, prior_exposure: float, responsivity: float = 1.0): + import numpy as np + + num_cameras, image_height, image_width, _ = hdr.shape + ds_height = max(1, image_height // PPISP_CONTROLLER_INPUT_DOWNSAMPLING) + ds_width = max(1, image_width // PPISP_CONTROLLER_INPUT_DOWNSAMPLING) + + conv1_w = weights[PPISP_CONTROLLER_OFF_CONV1_W:PPISP_CONTROLLER_OFF_CONV1_B].reshape(16, 3) + conv1_b = weights[PPISP_CONTROLLER_OFF_CONV1_B:PPISP_CONTROLLER_OFF_CONV2_W] + conv2_w = weights[PPISP_CONTROLLER_OFF_CONV2_W:PPISP_CONTROLLER_OFF_CONV2_B].reshape(32, 16) + conv2_b = weights[PPISP_CONTROLLER_OFF_CONV2_B:PPISP_CONTROLLER_OFF_CONV3_W] + conv3_w = weights[PPISP_CONTROLLER_OFF_CONV3_W:PPISP_CONTROLLER_OFF_CONV3_B].reshape(64, 32) + conv3_b = weights[PPISP_CONTROLLER_OFF_CONV3_B:PPISP_CONTROLLER_OFF_TRUNK0_W] + + pool1 = np.empty((num_cameras, ds_height, ds_width, 16), dtype=np.float32) + for camera_id in range(num_cameras): + for dy in range(ds_height): + y0 = dy * PPISP_CONTROLLER_INPUT_DOWNSAMPLING + y1 = min(y0 + PPISP_CONTROLLER_INPUT_DOWNSAMPLING, image_height) + for dx in range(ds_width): + x0 = dx * PPISP_CONTROLLER_INPUT_DOWNSAMPLING + x1 = min(x0 + PPISP_CONTROLLER_INPUT_DOWNSAMPLING, image_width) + conv1 = hdr[camera_id, y0:y1, x0:x1].reshape(-1, 3) * responsivity + conv1 = conv1 @ conv1_w.T + conv1_b + pool1[camera_id, dy, dx] = np.maximum(conv1.max(axis=0), 0.0) + + conv2 = np.maximum(np.einsum("nhwi,oi->nhwo", pool1, conv2_w) + conv2_b, 0.0) + conv3 = np.einsum("nhwi,oi->nhwo", conv2, conv3_w) + conv3_b + + pooled = np.zeros((num_cameras, 64, 5, 5), dtype=np.float32) + for gy in range(5): + h_start = (gy * ds_height) // 5 + h_end = min(((gy + 1) * ds_height + 4) // 5, ds_height) + for gx in range(5): + w_start = (gx * ds_width) // 5 + w_end = min(((gx + 1) * ds_width + 4) // 5, ds_width) + pooled[:, :, gy, gx] = conv3[:, h_start:h_end, w_start:w_end, :].mean(axis=(1, 2)) + + features = pooled.reshape(num_cameras, PPISP_CONTROLLER_FEATURE_LEN) + x = np.concatenate( + [features, np.full((num_cameras, 1), prior_exposure, dtype=np.float32)], + axis=1, + ) + + trunk0_w = weights[PPISP_CONTROLLER_OFF_TRUNK0_W:PPISP_CONTROLLER_OFF_TRUNK0_B].reshape( + PPISP_CONTROLLER_HIDDEN_DIM, PPISP_CONTROLLER_MLP_INPUT_DIM + ) + trunk0_b = weights[PPISP_CONTROLLER_OFF_TRUNK0_B:PPISP_CONTROLLER_OFF_TRUNK1_W] + trunk1_w = weights[PPISP_CONTROLLER_OFF_TRUNK1_W:PPISP_CONTROLLER_OFF_TRUNK1_B].reshape( + PPISP_CONTROLLER_HIDDEN_DIM, PPISP_CONTROLLER_HIDDEN_DIM + ) + trunk1_b = weights[PPISP_CONTROLLER_OFF_TRUNK1_B:PPISP_CONTROLLER_OFF_TRUNK2_W] + trunk2_w = weights[PPISP_CONTROLLER_OFF_TRUNK2_W:PPISP_CONTROLLER_OFF_TRUNK2_B].reshape( + PPISP_CONTROLLER_HIDDEN_DIM, PPISP_CONTROLLER_HIDDEN_DIM + ) + trunk2_b = weights[PPISP_CONTROLLER_OFF_TRUNK2_B:PPISP_CONTROLLER_OFF_EXP_W] + exp_w = weights[PPISP_CONTROLLER_OFF_EXP_W:PPISP_CONTROLLER_OFF_EXP_B] + exp_b = weights[PPISP_CONTROLLER_OFF_EXP_B] + col_w = weights[PPISP_CONTROLLER_OFF_COL_W:PPISP_CONTROLLER_OFF_COL_B].reshape(8, 128) + col_b = weights[PPISP_CONTROLLER_OFF_COL_B:PPISP_CONTROLLER_EXPECTED_WEIGHTS_LEN] + + hidden = np.maximum(x @ trunk0_w.T + trunk0_b, 0.0) + hidden = np.maximum(hidden @ trunk1_w.T + trunk1_b, 0.0) + hidden = np.maximum(hidden @ trunk2_w.T + trunk2_b, 0.0) + exposure = hidden @ exp_w + exp_b + color = hidden @ col_w.T + col_b + return np.concatenate([exposure[:, None], color], axis=1).astype(np.float32) + + def test_ppisp_warp_exposure_increases_ldr_output(): hdr_color = _hdr(0.25, (1, 4, 4, 3)) baseline = _rgba((1, 4, 4, 4)) @@ -50,6 +150,118 @@ def test_ppisp_warp_responsivity_lifts_dim_input(): assert boosted_np[..., :3].astype(float).mean() > baseline_np[..., :3].astype(float).mean() +def test_ppisp_warp_controller_exposure_increases_ldr_output(): + import numpy as np + + hdr_color = _hdr(0.25, (1, 4, 4, 3)) + baseline = _rgba((1, 4, 4, 4)) + exposed = _rgba((1, 4, 4, 4)) + baseline_params = wp.zeros((1, 9), dtype=wp.float32) + exposed_params = wp.from_numpy( + np.array([[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]], dtype=np.float32), + dtype=wp.float32, + ) + cfg = normalize_ppisp_cfg(PpispCfg()) + + apply_ppisp_to_rgba_with_controller_params(hdr_color, baseline, cfg, baseline_params) + apply_ppisp_to_rgba_with_controller_params(hdr_color, exposed, cfg, exposed_params) + + baseline_np = baseline.numpy() + exposed_np = exposed.numpy() + assert (baseline_np[..., 3] == 255).all() + assert (exposed_np[..., 3] == 255).all() + assert exposed_np[..., :3].astype(float).mean() > baseline_np[..., :3].astype(float).mean() + + +def test_ppisp_warp_controller_params_match_static_color_latent_order(): + import numpy as np + + hdr_data = np.linspace(0.05, 0.95, 1 * 4 * 4 * 3, dtype=np.float32).reshape(1, 4, 4, 3) + hdr_color = wp.from_numpy(hdr_data, dtype=wp.float32) + static = _rgba((1, 4, 4, 4)) + controller = _rgba((1, 4, 4, 4)) + inputs = { + "exposureOffset": 0.25, + "colorLatentBlue": (0.20, -0.10), + "colorLatentRed": (-0.15, 0.25), + "colorLatentGreen": (0.05, -0.30), + "colorLatentNeutral": (0.12, 0.08), + } + controller_params = wp.from_numpy( + np.array( + [ + [ + inputs["exposureOffset"], + *inputs["colorLatentBlue"], + *inputs["colorLatentRed"], + *inputs["colorLatentGreen"], + *inputs["colorLatentNeutral"], + ] + ], + dtype=np.float32, + ), + dtype=wp.float32, + ) + + apply_ppisp_to_rgba(hdr_color, static, normalize_ppisp_cfg(PpispCfg(inputs=inputs))) + apply_ppisp_to_rgba_with_controller_params( + hdr_color, + controller, + normalize_ppisp_cfg(PpispCfg()), + controller_params, + ) + + np.testing.assert_array_equal(controller.numpy(), static.numpy()) + + +@pytest.mark.parametrize("num_cameras", [1, 2]) +@pytest.mark.parametrize("image_shape", [(7, 8), (80, 96)]) +def test_ppisp_controller_network_matches_numpy_reference(num_cameras, image_shape): + import numpy as np + + if not wp.is_cuda_available(): + pytest.skip("PPISP controller requires CUDA.") + + rng = np.random.default_rng(7) + height, width = image_shape + hdr_np = rng.uniform(0.0, 1.0, size=(num_cameras, height, width, 3)).astype(np.float32) + weights_np = rng.normal(0.0, 0.003, size=PPISP_CONTROLLER_EXPECTED_WEIGHTS_LEN).astype(np.float32) + for offset, size in ( + (PPISP_CONTROLLER_OFF_CONV1_B, 16), + (PPISP_CONTROLLER_OFF_CONV2_B, 32), + (PPISP_CONTROLLER_OFF_CONV3_B, 64), + (PPISP_CONTROLLER_OFF_TRUNK0_B, PPISP_CONTROLLER_HIDDEN_DIM), + (PPISP_CONTROLLER_OFF_TRUNK1_B, PPISP_CONTROLLER_HIDDEN_DIM), + (PPISP_CONTROLLER_OFF_TRUNK2_B, PPISP_CONTROLLER_HIDDEN_DIM), + ): + weights_np[offset : offset + size] += 0.04 + prior_exposure = 0.35 + responsivity = 1.7 + + with wp.ScopedDevice("cuda:0"): + hdr_color = wp.from_numpy(hdr_np, dtype=wp.float32, device="cuda:0") + weights = wp.from_numpy(weights_np, dtype=wp.float32, device="cuda:0") + features = wp.empty((num_cameras, PPISP_CONTROLLER_FEATURE_LEN), dtype=wp.float32, device="cuda:0") + controller_params = wp.empty((num_cameras, PPISP_CONTROLLER_PARAM_COUNT), dtype=wp.float32, device="cuda:0") + + compute_ppisp_controller_params( + hdr_color, + weights, + features, + controller_params, + prior_exposure, + responsivity, + ) + wp.synchronize() + + np.testing.assert_allclose( + controller_params.numpy(), + _controller_reference_params(hdr_np, weights_np, prior_exposure, responsivity), + rtol=2.0e-4, + atol=2.0e-5, + ) + + def test_ppisp_warp_vignetting_uses_detiled_camera_coordinates(): hdr_color = _hdr(0.5, (2, 5, 5, 3)) rgba = _rgba((2, 5, 5, 4)) diff --git a/source/isaaclab_ppisp/test/test_ppisp_pipeline_perf.py b/source/isaaclab_ppisp/test/test_ppisp_pipeline_perf.py new file mode 100644 index 000000000000..29fe338bfe5a --- /dev/null +++ b/source/isaaclab_ppisp/test/test_ppisp_pipeline_perf.py @@ -0,0 +1,400 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Opt-in PPISP pipeline performance sweep. + +This test is skipped by default because it is a benchmarking sweep, not a +correctness test. The configured resolutions are final tiled render-target +resolutions; the benchmark derives each per-environment image size from the +same near-square layout used by tiled camera visualization. Run manually with: + +``` +./isaaclab.sh -p -m pip install pytest +ISAACLAB_PPISP_PERF=1 ./isaaclab.sh -p -m pytest \ + source/isaaclab_ppisp/test/test_ppisp_pipeline_perf.py -s +``` +""" + +from __future__ import annotations + +import csv +import gc +import math +import os +import re +from functools import lru_cache +from pathlib import Path + +import numpy as np +import pytest +import warp as wp +from isaaclab_ppisp import PPISP_CONTROLLER_EXPECTED_WEIGHTS_LEN, PpispCfg, PpispPipeline +from isaaclab_ppisp.kernels import ( + PPISP_CONTROLLER_FEATURE_LEN, + PPISP_CONTROLLER_HIDDEN_DIM, + PPISP_CONTROLLER_MLP_THREAD_GROUP_SIZE, + PPISP_CONTROLLER_OFF_CONV1_B, + PPISP_CONTROLLER_OFF_CONV2_B, + PPISP_CONTROLLER_OFF_CONV3_B, + PPISP_CONTROLLER_OFF_EXP_B, + PPISP_CONTROLLER_OFF_TRUNK0_B, + PPISP_CONTROLLER_OFF_TRUNK1_B, + PPISP_CONTROLLER_OFF_TRUNK2_B, + PPISP_CONTROLLER_PARAM_COUNT, + PPISP_CONTROLLER_POOL_CELL_COUNT, + PPISP_CONTROLLER_POOL_THREAD_GROUP_SIZE, + _ppisp_controller_mlp_native_kernel, + _ppisp_controller_pool_features_native_kernel, + apply_ppisp_to_rgba, + apply_ppisp_to_rgba_with_controller_params, +) + +pytestmark = pytest.mark.skipif( + os.environ.get("ISAACLAB_PPISP_PERF", "0") != "1", + reason="PPISP performance sweep is opt-in; set ISAACLAB_PPISP_PERF=1 to run.", +) + +_DEFAULT_RESOLUTIONS = ((256, 256), (512, 512), (1024, 1024), (2048, 2048)) +_DEFAULT_NUM_ENVS = (1, 8, 32, 128, 512) +_DEFAULT_WARMUP_ITERS = 2 +_DEFAULT_MEASURE_ITERS = 5 +_DEFAULT_MEMORY_FRACTION = 0.85 + + +def _parse_resolutions(value: str | None) -> tuple[tuple[int, int], ...]: + if not value: + return _DEFAULT_RESOLUTIONS + resolutions = [] + for item in value.split(","): + item = item.strip().lower() + if not item: + continue + if "x" in item: + height, width = item.split("x", maxsplit=1) + resolutions.append((int(height), int(width))) + else: + size = int(item) + resolutions.append((size, size)) + return tuple(resolutions) + + +def _parse_ints(value: str | None, default: tuple[int, ...]) -> tuple[int, ...]: + if not value: + return default + return tuple(int(item.strip()) for item in value.split(",") if item.strip()) + + +def _parse_variants(value: str | None) -> tuple[str, ...]: + if not value: + return ("static", "controller") + variants = tuple(item.strip().lower() for item in value.split(",") if item.strip()) + valid = {"static", "controller"} + unknown = set(variants) - valid + if unknown: + raise ValueError(f"Unknown PPISP perf variants: {sorted(unknown)}. Expected any of {sorted(valid)}.") + return variants + + +def _tile_grid_shape(num_envs: int) -> tuple[int, int]: + cols = max(1, math.ceil(math.sqrt(max(1, num_envs)))) + rows = math.ceil(max(1, num_envs) / cols) + return rows, cols + + +def _per_env_resolution(final_height: int, final_width: int, num_envs: int) -> tuple[int, int, int, int]: + tile_rows, tile_cols = _tile_grid_shape(num_envs) + return max(1, final_height // tile_rows), max(1, final_width // tile_cols), tile_rows, tile_cols + + +def _estimated_bytes(num_envs: int, height: int, width: int, *, controller: bool) -> int: + hdr_bytes = num_envs * height * width * 3 * np.dtype(np.float32).itemsize + rgba_bytes = num_envs * height * width * 4 * np.dtype(np.uint8).itemsize + controller_bytes = 0 + if controller: + controller_bytes += PPISP_CONTROLLER_EXPECTED_WEIGHTS_LEN * np.dtype(np.float32).itemsize + controller_bytes += num_envs * (1600 + 9) * np.dtype(np.float32).itemsize + return hdr_bytes + rgba_bytes + controller_bytes + + +@lru_cache(maxsize=1) +def _synthetic_controller_weights() -> tuple[float, ...]: + rng = np.random.default_rng(11) + weights = rng.normal(0.0, 0.003, size=PPISP_CONTROLLER_EXPECTED_WEIGHTS_LEN).astype(np.float32) + for offset, size in ( + (PPISP_CONTROLLER_OFF_CONV1_B, 16), + (PPISP_CONTROLLER_OFF_CONV2_B, 32), + (PPISP_CONTROLLER_OFF_CONV3_B, 64), + (PPISP_CONTROLLER_OFF_TRUNK0_B, PPISP_CONTROLLER_HIDDEN_DIM), + (PPISP_CONTROLLER_OFF_TRUNK1_B, PPISP_CONTROLLER_HIDDEN_DIM), + (PPISP_CONTROLLER_OFF_TRUNK2_B, PPISP_CONTROLLER_HIDDEN_DIM), + ): + weights[offset : offset + size] += 0.04 + weights[PPISP_CONTROLLER_OFF_EXP_B] = 0.0 + return tuple(float(value) for value in weights) + + +def _is_controller_variant(variant: str) -> bool: + return variant != "static" + + +def _make_pipeline(variant: str) -> PpispPipeline: + if variant == "static": + return PpispPipeline(PpispCfg(inputs={"exposureOffset": 0.0})) + return PpispPipeline(PpispCfg(controller_weights=_synthetic_controller_weights())) + + +def _time_pipeline_apply( + pipeline: PpispPipeline, + hdr: wp.array, + rgba: wp.array, + *, + warmup_iters: int, + measure_iters: int, +) -> float: + for _ in range(warmup_iters): + pipeline.apply(hdr, rgba) + wp.synchronize() + + start = wp.Event(enable_timing=True) + end = wp.Event(enable_timing=True) + wp.record_event(start) + for _ in range(measure_iters): + pipeline.apply(hdr, rgba) + wp.record_event(end) + return float(wp.get_event_elapsed_time(start, end)) / float(measure_iters) + + +def _time_gpu_operation(operation, *, warmup_iters: int, measure_iters: int) -> float: + for _ in range(warmup_iters): + operation() + wp.synchronize() + + start = wp.Event(enable_timing=True) + end = wp.Event(enable_timing=True) + wp.record_event(start) + for _ in range(measure_iters): + operation() + wp.record_event(end) + return float(wp.get_event_elapsed_time(start, end)) / float(measure_iters) + + +def _launch_controller_pool( + hdr: wp.array, + weights: wp.array, + features: wp.array, + responsivity: float, +) -> None: + wp.launch_tiled( + _ppisp_controller_pool_features_native_kernel, + dim=(int(hdr.shape[0]), PPISP_CONTROLLER_POOL_CELL_COUNT), + inputs=[hdr, weights, features, int(hdr.shape[2]), int(hdr.shape[1]), float(responsivity)], + device=str(hdr.device), + block_dim=PPISP_CONTROLLER_POOL_THREAD_GROUP_SIZE, + ) + + +def _launch_controller_mlp( + features: wp.array, + weights: wp.array, + controller_params: wp.array, + prior_exposure: float, +) -> None: + wp.launch_tiled( + _ppisp_controller_mlp_native_kernel, + dim=(int(features.shape[0]),), + inputs=[features, weights, controller_params, float(prior_exposure)], + device=str(features.device), + block_dim=PPISP_CONTROLLER_MLP_THREAD_GROUP_SIZE, + ) + + +def _print_perf_table(rows: list[dict[str, object]]) -> None: + print( + "\n[ppisp pipeline perf]\n" + "variant,num_envs,final_height,final_width,tile_rows,tile_cols,per_env_height,per_env_width," + "processed_mpix,final_mpix,estimated_gib,mean_ms,staged_total_ms,controller_pool_ms," + "controller_mlp_ms,apply_ms,stage_delta_ms,ms_per_env,throughput_mpix_s,status,error" + ) + for row in rows: + print( + "{variant},{num_envs},{final_height},{final_width},{tile_rows},{tile_cols}," + "{per_env_height},{per_env_width},{processed_mpix:.3f},{final_mpix:.3f},{estimated_gib:.2f}," + "{mean_ms:.3f},{staged_total_ms:.3f},{controller_pool_ms:.3f},{controller_mlp_ms:.3f}," + "{apply_ms:.3f},{stage_delta_ms:.3f},{ms_per_env:.5f},{throughput_mpix_s:.2f},{status},{error}".format( + **row + ) + ) + + +def _is_allocation_error(exc: Exception) -> bool: + message = str(exc).lower() + return bool(re.search(r"out[ -]of[ -]memory|failed to allocate|allocation", message)) + + +def test_ppisp_pipeline_perf_sweep(tmp_path): + if not wp.is_cuda_available(): + pytest.skip("PPISP performance sweep requires CUDA.") + + device_name = os.environ.get("ISAACLAB_PPISP_PERF_DEVICE", "cuda:0") + device = wp.get_device(device_name) + if not device.is_cuda: + pytest.skip(f"PPISP performance sweep requires a CUDA device, got {device_name}.") + + resolutions = _parse_resolutions(os.environ.get("ISAACLAB_PPISP_PERF_RESOLUTIONS")) + num_envs_values = _parse_ints(os.environ.get("ISAACLAB_PPISP_PERF_NUM_ENVS"), _DEFAULT_NUM_ENVS) + variants = _parse_variants(os.environ.get("ISAACLAB_PPISP_PERF_VARIANTS")) + warmup_iters = int(os.environ.get("ISAACLAB_PPISP_PERF_WARMUP_ITERS", _DEFAULT_WARMUP_ITERS)) + measure_iters = int(os.environ.get("ISAACLAB_PPISP_PERF_MEASURE_ITERS", _DEFAULT_MEASURE_ITERS)) + memory_fraction = float(os.environ.get("ISAACLAB_PPISP_PERF_MEMORY_FRACTION", _DEFAULT_MEMORY_FRACTION)) + max_case_bytes = int(device.total_memory * memory_fraction) + output_path = Path(os.environ.get("ISAACLAB_PPISP_PERF_OUTPUT", tmp_path / "ppisp_pipeline_perf.csv")) + + rows: list[dict[str, object]] = [] + pipelines = {variant: _make_pipeline(variant) for variant in variants} + with wp.ScopedDevice(device): + for final_height, final_width in resolutions: + for num_envs in num_envs_values: + height, width, tile_rows, tile_cols = _per_env_resolution(final_height, final_width, num_envs) + processed_mpix = (num_envs * height * width) / 1.0e6 + for variant in variants: + estimated_bytes = _estimated_bytes( + num_envs, + height, + width, + controller=_is_controller_variant(variant), + ) + row = { + "variant": variant, + "num_envs": num_envs, + "final_height": final_height, + "final_width": final_width, + "tile_rows": tile_rows, + "tile_cols": tile_cols, + "per_env_height": height, + "per_env_width": width, + "processed_mpix": processed_mpix, + "final_mpix": (final_height * final_width) / 1.0e6, + "estimated_gib": estimated_bytes / 1024**3, + "mean_ms": 0.0, + "staged_total_ms": 0.0, + "controller_pool_ms": 0.0, + "controller_mlp_ms": 0.0, + "apply_ms": 0.0, + "stage_delta_ms": 0.0, + "ms_per_env": 0.0, + "throughput_mpix_s": 0.0, + "status": "ok", + "error": "", + } + if estimated_bytes > max_case_bytes: + row["status"] = "skipped_memory_guard" + rows.append(row) + continue + + hdr = None + rgba = None + weights = None + features = None + controller_params = None + try: + hdr = wp.full( + (num_envs, height, width, 3), + value=wp.float32(0.25), + dtype=wp.float32, + device=device, + ) + rgba = wp.empty((num_envs, height, width, 4), dtype=wp.uint8, device=device) + pipeline = pipelines[variant] + mean_ms = _time_pipeline_apply( + pipeline, + hdr, + rgba, + warmup_iters=warmup_iters, + measure_iters=measure_iters, + ) + if _is_controller_variant(variant): + weights = wp.array(pipeline.cfg.controller_weights, dtype=wp.float32, device=device) + features = wp.empty( + (num_envs, PPISP_CONTROLLER_FEATURE_LEN), dtype=wp.float32, device=device + ) + controller_params = wp.empty( + (num_envs, PPISP_CONTROLLER_PARAM_COUNT), dtype=wp.float32, device=device + ) + _launch_controller_pool(hdr, weights, features, float(pipeline.cfg.controller_responsivity)) + _launch_controller_mlp( + features, + weights, + controller_params, + pipeline.cfg.controller_prior_exposure, + ) + wp.synchronize() + controller_pool_ms = _time_gpu_operation( + lambda: _launch_controller_pool( + hdr, weights, features, float(pipeline.cfg.controller_responsivity) + ), + warmup_iters=warmup_iters, + measure_iters=measure_iters, + ) + controller_mlp_ms = _time_gpu_operation( + lambda: _launch_controller_mlp( + features, + weights, + controller_params, + pipeline.cfg.controller_prior_exposure, + ), + warmup_iters=warmup_iters, + measure_iters=measure_iters, + ) + apply_ms = _time_gpu_operation( + lambda: apply_ppisp_to_rgba_with_controller_params( + hdr, rgba, pipeline.cfg, controller_params + ), + warmup_iters=warmup_iters, + measure_iters=measure_iters, + ) + staged_total_ms = controller_pool_ms + controller_mlp_ms + apply_ms + row["controller_pool_ms"] = controller_pool_ms + row["controller_mlp_ms"] = controller_mlp_ms + else: + apply_ms = _time_gpu_operation( + lambda: apply_ppisp_to_rgba(hdr, rgba, pipeline.cfg), + warmup_iters=warmup_iters, + measure_iters=measure_iters, + ) + staged_total_ms = apply_ms + row["staged_total_ms"] = staged_total_ms + row["apply_ms"] = apply_ms + row["stage_delta_ms"] = mean_ms - staged_total_ms + row["mean_ms"] = mean_ms + row["ms_per_env"] = mean_ms / float(num_envs) + row["throughput_mpix_s"] = row["processed_mpix"] / (mean_ms / 1000.0) + except Exception as exc: + if _is_allocation_error(exc): + row["status"] = "skipped_allocation" + else: + row["status"] = f"failed:{type(exc).__name__}" + row["error"] = str(exc).replace("\n", " ")[:160] + finally: + hdr = None + rgba = None + weights = None + features = None + controller_params = None + gc.collect() + wp.synchronize() + rows.append(row) + + output_path.parent.mkdir(parents=True, exist_ok=True) + with output_path.open("w", newline="", encoding="utf-8") as csv_file: + writer = csv.DictWriter(csv_file, fieldnames=list(rows[0])) + writer.writeheader() + writer.writerows(rows) + + _print_perf_table(rows) + print(f"[ppisp pipeline perf] wrote {output_path}") + assert rows + assert any(row["status"] == "ok" for row in rows) + failed_rows = [row for row in rows if str(row["status"]).startswith("failed:")] + assert not failed_rows