diff --git a/compass/residual_rl/residual_ppo_trainer.py b/compass/residual_rl/residual_ppo_trainer.py index 36447cc..af5da82 100644 --- a/compass/residual_rl/residual_ppo_trainer.py +++ b/compass/residual_rl/residual_ppo_trainer.py @@ -54,7 +54,8 @@ def __init__(self, debug_viz=False, max_debug_images=2, debug_image_interval=10, - debug_image_step=0): + debug_image_step=0, + save_debug_viewport_images=False): # Prepare log directory. exist_ok=True avoids a TOCTOU race when several # torchrun ranks check + create concurrently. os.makedirs(output_dir, exist_ok=True) @@ -128,6 +129,7 @@ def __init__(self, # Save one rollout step every `debug_image_interval` iterations. self.debug_image_interval = debug_image_interval self.debug_image_step = debug_image_step + self.save_debug_viewport_images = save_debug_viewport_images self.env.reset() @@ -583,9 +585,9 @@ def _upload_video(self, iteration): # Video generation can be slower than this function call, so add one iteration # delay for video upload. target_iteration = iteration - 1 - target_video_path = os.path.join( - self.output_dir, 'videos', - f'rl-video-step-{target_iteration*self.num_steps_per_env}.mp4') + video_step = target_iteration * self.num_steps_per_env + target_video_path = os.path.join(self.output_dir, 'videos', 'robot_camera', + f'rl-video-step-{video_step}.mp4') print(target_video_path) if os.path.exists(target_video_path): self.logger.log_video(name="episode_video", @@ -726,7 +728,8 @@ def _save_debug_images(self, obs_dict, iteration, step): image_path=grid_image_path, step=iteration) - self._save_debug_viewport_image(iteration, step) + if self.save_debug_viewport_images: + self._save_debug_viewport_image(iteration, step) except (KeyError, IOError, OSError, ValueError, RuntimeError, AttributeError) as e: print(f"Warning: Failed to save debug images: {e}") diff --git a/compass/rl_env/exts/mobility_es/mobility_es/config/nurec_scenes.py b/compass/rl_env/exts/mobility_es/mobility_es/config/nurec_scenes.py index e0f3e22..0b4a29e 100644 --- a/compass/rl_env/exts/mobility_es/mobility_es/config/nurec_scenes.py +++ b/compass/rl_env/exts/mobility_es/mobility_es/config/nurec_scenes.py @@ -20,7 +20,6 @@ import isaaclab.sim as sim_utils from isaaclab.assets import AssetBaseCfg -from compass.utils.nurec_utils import NUREC_SCENE_NAMES from mobility_es.config.environments import EnvSceneAssetCfg, OMAP_PATHS, USD_PATHS _USD_DIR = os.path.join(os.path.dirname(__file__), "../usd") @@ -50,15 +49,29 @@ def _make_prim_leaf(scene_name: str) -> str: "omap_file": "occupancy_map_with_sim_objects.yaml", }, } -NUREC_SCENES = [ - NurecScene(scene_name, _make_prim_leaf(scene_name), **NUREC_SCENE_OVERRIDES.get(scene_name, {})) - for scene_name in NUREC_SCENE_NAMES -] -def _make_nurec_scene_asset_cfg(scene: NurecScene, usd_file: str) -> EnvSceneAssetCfg: +def _make_nurec_scene(scene_name: str, omap_file: str | None = None) -> NurecScene: + """Build NuRec scene metadata from a scene folder name.""" + overrides = dict(NUREC_SCENE_OVERRIDES.get(scene_name, {})) + if omap_file is not None: + overrides["omap_file"] = omap_file + return NurecScene(scene_name, _make_prim_leaf(scene_name), **overrides) + + +def make_nurec_scene_asset_cfg(scene_name: str, + usd_file: str, + omap_file: str | None = None) -> EnvSceneAssetCfg: """Build the shared :class:`EnvSceneAssetCfg` for a NuRec Real2Sim scene.""" + scene = _make_nurec_scene(scene_name, omap_file) usd_dir = os.path.join(_USD_DIR, scene.folder) + usd_path = os.path.join(usd_dir, usd_file) + omap_path = os.path.join(usd_dir, scene.omap_file) + USD_PATHS[scene.prim_leaf] = usd_path + OMAP_PATHS[scene.prim_leaf] = { + "path": omap_path, + "origin_convention": scene.origin_convention, + } return EnvSceneAssetCfg( prim_path="{ENV_REGEX_NS}/" + scene.prim_leaf, init_state=AssetBaseCfg.InitialStateCfg( @@ -66,7 +79,7 @@ def _make_nurec_scene_asset_cfg(scene: NurecScene, usd_file: str) -> EnvSceneAss rot=(0.0, 0.0, 0.0, 1.0), ), spawn=sim_utils.UsdFileCfg( - usd_path=os.path.join(usd_dir, usd_file), + usd_path=usd_path, scale=(1.0, 1.0, 1.0), rigid_props=sim_utils.RigidBodyPropertiesCfg( disable_gravity=None, @@ -76,17 +89,3 @@ def _make_nurec_scene_asset_cfg(scene: NurecScene, usd_file: str) -> EnvSceneAss ), env_spacing=scene.env_spacing, ) - - -def make_nurec_scene_asset_cfg_map(usd_file: str) -> dict[str, EnvSceneAssetCfg]: - """Register NuRec paths and build scene configs keyed by ``--environment``.""" - envs = {} - for scene in NUREC_SCENES: - usd_dir = os.path.join(_USD_DIR, scene.folder) - USD_PATHS[scene.prim_leaf] = os.path.join(usd_dir, usd_file) - OMAP_PATHS[scene.prim_leaf] = { - "path": os.path.join(usd_dir, scene.omap_file), - "origin_convention": scene.origin_convention, - } - envs[scene.folder] = _make_nurec_scene_asset_cfg(scene, usd_file) - return envs diff --git a/compass/rl_env/exts/mobility_es/mobility_es/config/scene_assets.py b/compass/rl_env/exts/mobility_es/mobility_es/config/scene_assets.py index b8d8237..7dc7db0 100644 --- a/compass/rl_env/exts/mobility_es/mobility_es/config/scene_assets.py +++ b/compass/rl_env/exts/mobility_es/mobility_es/config/scene_assets.py @@ -18,7 +18,6 @@ from isaaclab.sensors import CameraCfg, TiledCameraCfg from isaaclab.sensors.camera.camera_isp import CameraISPMode from isaaclab.terrains import TerrainImporterCfg -from isaaclab_physx.renderers import IsaacRtxRendererCfg, IsaacRtxRendererGlobalSettingsCfg # Terrain terrain = TerrainImporterCfg( @@ -41,16 +40,6 @@ width=512, data_types=["rgb", "depth"], isp_cfg=CameraISPMode.AUTO_ANY, - renderer_cfg=IsaacRtxRendererCfg(global_settings=IsaacRtxRendererGlobalSettingsCfg( - enable_dl_denoiser=True, - antialiasing_mode="DLSS", - carb_settings={ - "/rtx/rtpt/gaussian/skipTonemapping/enabled": False, - "/rtx/rtpt/gaussian/accumulatedDepth/allHits/enabled": True, - "/rtx/rtpt/gaussian/accumulatedAlbedo/enabled": True, - "/rtx/rtpt/gaussian/maxGaussiansToAccumulate": 360, - }, - )), spawn=sim_utils.PinholeCameraCfg(focal_length=10.0, focus_distance=400.0, horizontal_aperture=20.955, diff --git a/compass/utils/logger.py b/compass/utils/logger.py index b004742..c8f397a 100644 --- a/compass/utils/logger.py +++ b/compass/utils/logger.py @@ -17,11 +17,14 @@ """ import os +import re from typing import Dict, Any, Optional import wandb from tensorboardX import SummaryWriter +_ARTIFACT_COMPONENT_PATTERN = re.compile(r"[^A-Za-z0-9_.-]+") + class Logger: """ @@ -70,6 +73,16 @@ def __init__(self, print(f"Logger initialized with backend: {self.backend}") + @staticmethod + def _sanitize_artifact_component(value: str) -> str: + """Convert a run name into a W&B artifact-name-safe component.""" + sanitized = _ARTIFACT_COMPONENT_PATTERN.sub("-", value).strip("._-") + return sanitized or "run" + + def _run_scoped_artifact_name(self, name: str) -> str: + run_name = self._sanitize_artifact_component(self.experiment_name) + return f"{name}_{run_name}" + def log_image(self, name: str, image_path: str, step: Optional[int] = None): """ Log an image. @@ -136,13 +149,14 @@ def log_artifact(self, type: Type of the artifact description: Description of the artifact """ + artifact_name = self._run_scoped_artifact_name(name) if self.backend == "wandb": - artifact = wandb.Artifact(name=name, type=type, description=description) + artifact = wandb.Artifact(name=artifact_name, type=type, description=description) artifact.add_file(artifact_path) wandb.log_artifact(artifact) else: # TensorBoard doesn't support artifacts directly - print(f"Artifacts not supported in TensorBoard. Artifact '{name}' not logged.") + print(f"Artifacts not supported in TensorBoard. Artifact '{artifact_name}' not logged.") def log_config(self, config: Dict[str, Any]): """ diff --git a/compass/utils/nurec_utils.py b/compass/utils/nurec_utils.py index 6e4bc65..0900701 100644 --- a/compass/utils/nurec_utils.py +++ b/compass/utils/nurec_utils.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -15,329 +15,20 @@ import os -# pylint: disable=import-outside-toplevel - -NUREC_SCENE_NAMES = [ - "nova_carter-galileo", - "nova_carter-cafe", - "hand_hold-endeavor-andoria", - "hand_hold-endeavor-livingroom", - "hand_hold-endeavor-wormhole", - "hand_hold-endeavor-wormhole-table", - "hand_hold-voyager-babyboom", - "xgrid-wormhole", -] +# NuRec Asset Selection PARTICLE_SPG_RUNTIME_USD_FILE = "particle_spg-runtime.usdz" -NUREC_SPG_RUNTIME_KIT_ARGS = ( - "--/rtx/spg/enabled=true", - "--/omni/rtx/nre/compositing/disableNuRecPostProcessings=true", - "--/rtx/rtpt/gaussian/skipTonemapping/enabled=false", - "--enable omni.rtx.spg", -) -KIT_PERSPECTIVE_CAMERA_PATH = "/OmniverseKit_Persp" -KIT_SCENE_PARTITION = "env_0" -KIT_PPISP_RENDER_PRODUCT_PATH = "/Render/COMPASS_KitPerspective_PPISP" -NUREC_IDENTITY_EXPOSURE = { - "exposure": 0.0, - "exposure:fStop": 1.0, - "exposure:iso": 0.0, - "exposure:responsivity": 1.0, - "exposure:time": 1.0, -} - -def apply_nurec_spg_kit_args(args): - """Prepend SPG Kit args when the selected NuRec USD requires SPG runtime.""" - args.spg_runtime = uses_nurec_spg_runtime(args) - if not args.spg_runtime: - return - - existing_kit_args = getattr(args, "kit_args", None) or "" - if isinstance(existing_kit_args, list): - existing_kit_args = " ".join(existing_kit_args) - additions = [arg for arg in NUREC_SPG_RUNTIME_KIT_ARGS if arg not in existing_kit_args] - args.kit_args = " ".join(additions + ([existing_kit_args] if existing_kit_args else [])) +__all__ = ["PARTICLE_SPG_RUNTIME_USD_FILE", "uses_nurec_spg_runtime"] +# Public API def uses_nurec_spg_runtime(args): """Return whether the CLI selection should enable NuRec SPG runtime.""" if getattr(args, "spg_runtime", False): return True - if getattr(args, "nurec_scene", None) not in NUREC_SCENE_NAMES: + if getattr(args, "nurec_scene", None) is None: return False usd_file = getattr(args, "nurec_usd_file", "") or "" return os.path.basename(usd_file) == PARTICLE_SPG_RUNTIME_USD_FILE - - -def _render_product_camera_path(render_product_prim): - """Return the first camera target from a RenderProduct relationship.""" - from pxr import UsdRender - - product = UsdRender.Product(render_product_prim) - camera_rel = render_product_prim.GetRelationship("camera") - targets = camera_rel.GetTargets() if camera_rel else [] - if not targets: - targets = product.GetCameraRel().GetForwardedTargets() - return targets[0] if targets else None - - -def _is_render_product_prim(prim): - """Return whether a prim is a USD RenderProduct.""" - from pxr import UsdRender - - return prim.GetTypeName() == "RenderProduct" or prim.IsA(UsdRender.Product) - - -def _iter_render_products(stage): - """Yield RenderProduct prims in stage traversal order.""" - for prim in stage.Traverse(): - if _is_render_product_prim(prim): - yield prim - - -def _format_render_product_choices(stage): - """Return RenderProduct choices in the same path-to-camera form used in warnings.""" - choices = [] - for prim in _iter_render_products(stage): - camera_path = _render_product_camera_path(prim) or "" - choices.append(f"{prim.GetPath()} -> {camera_path}") - return "\n".join(f" {choice}" for choice in sorted(choices)) or " " - - -def _find_source_render_product(stage): - """Return the first RenderProduct with a camera target, if any.""" - for prim in _iter_render_products(stage): - camera_path = _render_product_camera_path(prim) - if camera_path is not None: - return prim, camera_path - return None - - -def _get_source_render_product(stage, stage_label): - """Return the source RenderProduct or raise with available choices.""" - discovered = _find_source_render_product(stage) - if discovered is not None: - return discovered - - raise RuntimeError(f"Could not find NuRec PPISP RenderProduct in {stage_label}. " - f"Available RenderProducts:\n{_format_render_product_choices(stage)}") - - -def _package_qualified_asset(package_path, asset_path): - """Qualify package-internal asset references with their USD package.""" - if not asset_path or "[" in asset_path: - return asset_path - return f"{package_path}[{os.path.basename(asset_path)}]" - - -def _copy_source_ppisp_render_product(stage, source_stage, source_usd_path, kit_camera_path): - """Copy the source PPISP RenderProduct and retarget it to the Kit camera. - - This mirrors show-renderproduct.py's CopyPrim flow, but uses Sdf.CopySpec so - the RenderProduct can be copied from the source USD package even when it is - not present in the composed GUI stage. - """ - from pxr import Sdf, Usd, UsdRender - - source_rp_prim, source_camera_path = _get_source_render_product(source_stage, source_usd_path) - prim_stack = source_rp_prim.GetPrimStack() - if not prim_stack: - raise RuntimeError(f"PPISP RenderProduct has empty prim stack: {source_rp_prim.GetPath()}") - - source_layer = prim_stack[0].layer - session = stage.GetSessionLayer() - src_path = Sdf.Path(str(source_rp_prim.GetPath())) - dst_path = Sdf.Path(KIT_PPISP_RENDER_PRODUCT_PATH) - package_path = source_stage.GetRootLayer().identifier - - with Usd.EditContext(stage, session): - existing = stage.GetPrimAtPath(dst_path) - if existing and existing.IsValid(): - stage.RemovePrim(dst_path) - - Sdf.CreatePrimInLayer(session, dst_path) - Sdf.CopySpec(source_layer, src_path, session, dst_path) - - with Usd.EditContext(stage, session): - dst_prim = stage.GetPrimAtPath(dst_path) - for prim in Usd.PrimRange(dst_prim): - for attr in prim.GetAttributes(): - connections = attr.GetConnections() - remapped = [conn.ReplacePrefix(src_path, dst_path) for conn in connections] - if remapped != connections: - attr.SetConnections(remapped) - - value = attr.Get() - if isinstance(value, Sdf.AssetPath): - attr.Set(Sdf.AssetPath(_package_qualified_asset(package_path, value.path))) - - for rel in prim.GetRelationships(): - targets = rel.GetTargets() - remapped = [target.ReplacePrefix(src_path, dst_path) for target in targets] - if remapped != targets: - rel.SetTargets(remapped) - - spec = session.GetPrimAtPath(prim.GetPath()) - if spec and spec.referenceList.prependedItems: - spec.referenceList.prependedItems = [ - Sdf.Reference(_package_qualified_asset(package_path, ref.assetPath)) - for ref in spec.referenceList.prependedItems - ] - - UsdRender.Product(dst_prim).GetCameraRel().SetTargets([Sdf.Path(kit_camera_path)]) - - return KIT_PPISP_RENDER_PRODUCT_PATH, source_camera_path - - -def _apply_render_settings_from_stage(stage): - """Apply render settings authored in the source stage custom layer data.""" - import carb.settings - - layer_data = stage.GetRootLayer().customLayerData or {} - render_settings = layer_data.get("renderSettings") or {} - if not render_settings: - return - - settings = carb.settings.get_settings() - for key, value in render_settings.items(): - settings.set("/" + str(key).replace(":", "/"), value) - - -def _copy_camera_metadata_to_chase_camera(stage, source_stage, source_camera_path, kit_camera_path): - """Copy NuRec PPISP camera metadata while preserving the Kit chase-camera pose.""" - from pxr import Sdf, Usd - - source_camera_prim = source_stage.GetPrimAtPath(source_camera_path) - if (not source_camera_prim or not source_camera_prim.IsValid() - or source_camera_prim.GetTypeName() != "Camera"): - raise RuntimeError(f"Viewport source camera does not exist: {source_camera_path}") - - chase_camera_prim = stage.GetPrimAtPath(kit_camera_path) - if (not chase_camera_prim or not chase_camera_prim.IsValid() - or chase_camera_prim.GetTypeName() != "Camera"): - raise RuntimeError(f"Kit chase camera does not exist: {kit_camera_path}") - - with Usd.EditContext(stage, stage.GetSessionLayer()): - api_schemas = source_camera_prim.GetMetadata("apiSchemas") - if api_schemas: - chase_camera_prim.SetMetadata("apiSchemas", api_schemas) - - for src_attr in source_camera_prim.GetAuthoredAttributes(): - name = src_attr.GetName() - if not name.startswith("ppisp:"): - continue - - value = src_attr.Get() - if value is None: - continue - - dst_attr = chase_camera_prim.GetAttribute(name) - if not dst_attr: - dst_attr = chase_camera_prim.CreateAttribute(name, src_attr.GetTypeName()) - dst_attr.Set(value) - - chase_camera_prim.AddAppliedSchema("OmniRtxCameraAutoExposureAPI_1") - chase_camera_prim.AddAppliedSchema("OmniRtxCameraExposureAPI_1") - for name, value in NUREC_IDENTITY_EXPOSURE.items(): - chase_camera_prim.CreateAttribute(name, Sdf.ValueTypeNames.Float).Set(value) - chase_camera_prim.CreateAttribute("omni:rtx:autoExposure:enabled", - Sdf.ValueTypeNames.Bool).Set(False) - - -def _bind_nurec_ppisp_render_product(stage, viewport, nurec_usd_path, kit_camera_path, quiet=False): - """Bind the Kit viewport to a copied NuRec PPISP RenderProduct.""" - if not nurec_usd_path: - return None - - try: - from pxr import Sdf, Usd - - source_stage = Usd.Stage.Open(nurec_usd_path) - if source_stage is None: - raise RuntimeError(f"Could not open NuRec USD: {nurec_usd_path}") - - render_product_path, source_camera_path = _copy_source_ppisp_render_product( - stage, source_stage, nurec_usd_path, kit_camera_path) - _copy_camera_metadata_to_chase_camera(stage, source_stage, source_camera_path, - kit_camera_path) - _apply_render_settings_from_stage(source_stage) - - viewport.camera_path = Sdf.Path(kit_camera_path) - viewport.render_product_path = render_product_path - if str(getattr(viewport, "render_product_path", "")) != render_product_path: - if not quiet: - print(f"[WARN] Kit ignored viewport RenderProduct " - f"{render_product_path}; keeping default viewport render product.") - return None - - return render_product_path - except Exception as exc: # pylint: disable=broad-except - if not quiet: - print(f"[WARN] Could not bind Kit viewport RenderProduct: " - f"{type(exc).__name__}: {exc}") - return None - - -def configure_nurec_kit_viewport(nurec_usd_path, spg_runtime, quiet=False): - """Configure the active Kit viewport for NuRec scene visualization. - - The viewport keeps COMPASS's external chase camera. For SPG runtime assets, - copy the source PPISP RenderProduct and bind it to that camera so the - viewport uses the same sensor-processing path as show-renderproduct.py. - """ - try: - import omni.kit.viewport.utility as viewport_utils - import omni.usd - from pxr import Sdf - except ImportError as exc: - if not quiet: - print(f"[WARN] Could not import USD utilities for Kit camera setup: {exc}") - return - - stage = omni.usd.get_context().get_stage() - if stage is None: - if not quiet: - print("[WARN] Could not configure Kit camera: no active USD stage.") - return - - camera_prim = stage.GetPrimAtPath(KIT_PERSPECTIVE_CAMERA_PATH) - if not camera_prim or not camera_prim.IsValid(): - if not quiet: - print(f"[WARN] Could not configure Kit camera: " - f"{KIT_PERSPECTIVE_CAMERA_PATH} does not exist.") - return - - attr = camera_prim.GetAttribute("omni:scenePartition") - if not attr.IsValid(): - attr = camera_prim.CreateAttribute("omni:scenePartition", Sdf.ValueTypeNames.Token) - attr.Set(KIT_SCENE_PARTITION) - - render_product_path = None - viewport = viewport_utils.get_active_viewport() - if viewport is not None: - try: - viewport.camera_path = Sdf.Path(KIT_PERSPECTIVE_CAMERA_PATH) - except Exception: # pylint: disable=broad-except - pass - if spg_runtime: - bound_nurec_rp_path = _bind_nurec_ppisp_render_product(stage, - viewport, - nurec_usd_path, - KIT_PERSPECTIVE_CAMERA_PATH, - quiet=quiet) - else: - bound_nurec_rp_path = None - render_product_path = getattr(viewport, "render_product_path", None) - else: - bound_nurec_rp_path = None - - if not quiet: - message = (f"[INFO] Set {KIT_PERSPECTIVE_CAMERA_PATH} scene partition to " - f"{KIT_SCENE_PARTITION}.") - if bound_nurec_rp_path: - message += f" Bound Kit viewport RenderProduct: {bound_nurec_rp_path}." - if render_product_path: - message += f" Active Kit render product: {render_product_path}." - print(message) diff --git a/compass/utils/renderer_utils.py b/compass/utils/renderer_utils.py new file mode 100644 index 0000000..d79bf58 --- /dev/null +++ b/compass/utils/renderer_utils.py @@ -0,0 +1,387 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os + +from compass.utils.nurec_utils import uses_nurec_spg_runtime +from compass.utils.visualizer_utils import ( + KIT_PERSPECTIVE_CAMERA_PATH, + KIT_SCENE_PARTITION, +) + +# pylint: disable=import-outside-toplevel + +############################### +# Isaac Sim RTX Renderer Setup +############################### +__all__ = [ + "apply_isaac_rtx_camera_renderer_settings", + "apply_nurec_spg_kit_args", + "configure_nurec_isaacsim_rtx_viewport", +] + +# Runtime args when launching Isaac Sim with a NuRec USD that requires SPG runtime. +NUREC_SPG_RUNTIME_KIT_ARGS = ( + "--/rtx/spg/enabled=true", + "--/omni/rtx/nre/compositing/disableNuRecPostProcessings=true", + "--enable omni.rtx.spg", +) + +# Camera renderer settings for Isaac RTX. +ISAAC_RTX_CAMERA_RENDERER_SETTINGS = { + "enable_dl_denoiser": True, + "antialiasing_mode": "DLSS", +} + +# Additional camera renderer settings for NuRec USDs that require SPG runtime. +NUREC_SPG_CAMERA_RENDERER_CARB_SETTINGS = { + "/rtx/rtpt/gaussian/skipTonemapping/enabled": False, + "/rtx/rtpt/gaussian/accumulatedDepth/allHits/enabled": True, + "/rtx/rtpt/gaussian/accumulatedAlbedo/enabled": True, + "/rtx/rtpt/gaussian/maxGaussiansToAccumulate": 360, +} + +# RenderProduct defaults for kit viewport when using NuRec USDs with SPG runtime. +KIT_PPISP_RENDER_PRODUCT_PATH = "/Render/COMPASS_KitPerspective_PPISP" +NUREC_IDENTITY_EXPOSURE = { + "exposure": 0.0, + "exposure:fStop": 1.0, + "exposure:iso": 0.0, + "exposure:responsivity": 1.0, + "exposure:time": 1.0, +} + + +# Public API +def apply_isaac_rtx_camera_renderer_settings(env_cfg, isaac_rtx, spg_runtime): + """Apply Isaac RTX camera renderer settings when the Isaac RTX renderer is used.""" + if not isaac_rtx: + return + + camera_cfg = getattr(env_cfg.scene, "camera", None) + if camera_cfg is None: + return + + from isaaclab_physx.renderers import ( + IsaacRtxRendererCfg, + IsaacRtxRendererGlobalSettingsCfg, + ) + + renderer_cfg = getattr(camera_cfg, "renderer_cfg", None) + if getattr(renderer_cfg, "renderer_type", None) != "isaac_rtx": + renderer_cfg = IsaacRtxRendererCfg(global_settings=IsaacRtxRendererGlobalSettingsCfg()) + camera_cfg.renderer_cfg = renderer_cfg + + global_settings = getattr(renderer_cfg, "global_settings", None) + if global_settings is None: + return + + for name, value in ISAAC_RTX_CAMERA_RENDERER_SETTINGS.items(): + setattr(global_settings, name, value) + + if not spg_runtime: + return + + carb_settings = dict(getattr(global_settings, "carb_settings", None) or {}) + carb_settings.update(NUREC_SPG_CAMERA_RENDERER_CARB_SETTINGS) + global_settings.carb_settings = carb_settings + + +def apply_nurec_spg_kit_args(args): + """Prepend SPG Kit args when the selected NuRec USD requires SPG runtime.""" + args.spg_runtime = uses_nurec_spg_runtime(args) + if not args.spg_runtime: + return + + existing_kit_args = getattr(args, "kit_args", None) or "" + if isinstance(existing_kit_args, list): + existing_kit_args = " ".join(existing_kit_args) + additions = [arg for arg in NUREC_SPG_RUNTIME_KIT_ARGS if arg not in existing_kit_args] + args.kit_args = " ".join(additions + ([existing_kit_args] if existing_kit_args else [])) + + +def configure_nurec_isaacsim_rtx_viewport(nurec_usd_path, spg_runtime, quiet=False): + """Configure the active Kit viewport for NuRec scene visualization. + + The viewport keeps COMPASS's external chase camera. For SPG runtime assets, + copy the source PPISP RenderProduct and bind it to that camera so the + viewport uses the same sensor-processing path as show-renderproduct.py. + """ + try: + import omni.kit.viewport.utility as viewport_utils + import omni.usd + from pxr import Sdf + except ImportError as exc: + if not quiet: + print(f"[WARN] Could not import USD utilities for Kit camera setup: {exc}") + return + + stage = omni.usd.get_context().get_stage() + if stage is None: + if not quiet: + print("[WARN] Could not configure Kit camera: no active USD stage.") + return + + camera_prim = stage.GetPrimAtPath(KIT_PERSPECTIVE_CAMERA_PATH) + if not camera_prim or not camera_prim.IsValid(): + if not quiet: + print(f"[WARN] Could not configure Kit camera: " + f"{KIT_PERSPECTIVE_CAMERA_PATH} does not exist.") + return + + attr = camera_prim.GetAttribute("omni:scenePartition") + if not attr.IsValid(): + attr = camera_prim.CreateAttribute("omni:scenePartition", Sdf.ValueTypeNames.Token) + attr.Set(KIT_SCENE_PARTITION) + + render_product_path = None + viewport = viewport_utils.get_active_viewport() + if viewport is not None: + try: + viewport.camera_path = Sdf.Path(KIT_PERSPECTIVE_CAMERA_PATH) + except Exception: # pylint: disable=broad-except + pass + if spg_runtime: + bound_nurec_rp_path = _bind_nurec_ppisp_render_product( + stage, + viewport, + nurec_usd_path, + KIT_PERSPECTIVE_CAMERA_PATH, + quiet=quiet, + ) + else: + bound_nurec_rp_path = None + render_product_path = getattr(viewport, "render_product_path", None) + else: + bound_nurec_rp_path = None + + if not quiet: + message = (f"[INFO] Set {KIT_PERSPECTIVE_CAMERA_PATH} scene partition to " + f"{KIT_SCENE_PARTITION}.") + if bound_nurec_rp_path: + message += f" Bound Kit viewport RenderProduct: {bound_nurec_rp_path}." + if render_product_path: + message += f" Active Kit render product: {render_product_path}." + print(message) + + +# Private helpers +def _render_product_camera_path(render_product_prim): + """Return the first camera target from a RenderProduct relationship.""" + from pxr import UsdRender + + product = UsdRender.Product(render_product_prim) + camera_rel = render_product_prim.GetRelationship("camera") + targets = camera_rel.GetTargets() if camera_rel else [] + if not targets: + targets = product.GetCameraRel().GetForwardedTargets() + return targets[0] if targets else None + + +def _is_render_product_prim(prim): + """Return whether a prim is a USD RenderProduct.""" + from pxr import UsdRender + + return prim.GetTypeName() == "RenderProduct" or prim.IsA(UsdRender.Product) + + +def _iter_render_products(stage): + """Yield RenderProduct prims in stage traversal order.""" + for prim in stage.Traverse(): + if _is_render_product_prim(prim): + yield prim + + +def _format_render_product_choices(stage): + """Return RenderProduct choices in the same path-to-camera form used in warnings.""" + choices = [] + for prim in _iter_render_products(stage): + camera_path = _render_product_camera_path(prim) or "" + choices.append(f"{prim.GetPath()} -> {camera_path}") + return "\n".join(f" {choice}" for choice in sorted(choices)) or " " + + +def _find_source_render_product(stage): + """Return the first RenderProduct with a camera target, if any.""" + for prim in _iter_render_products(stage): + camera_path = _render_product_camera_path(prim) + if camera_path is not None: + return prim, camera_path + return None + + +def _get_source_render_product(stage, stage_label): + """Return the source RenderProduct or raise with available choices.""" + discovered = _find_source_render_product(stage) + if discovered is not None: + return discovered + + raise RuntimeError(f"Could not find NuRec PPISP RenderProduct in {stage_label}. " + f"Available RenderProducts:\n{_format_render_product_choices(stage)}") + + +def _package_qualified_asset(package_path, asset_path): + """Qualify package-internal asset references with their USD package.""" + if not asset_path or "[" in asset_path: + return asset_path + return f"{package_path}[{os.path.basename(asset_path)}]" + + +def _copy_source_ppisp_render_product(stage, source_stage, source_usd_path, kit_camera_path): + """Copy the source PPISP RenderProduct and retarget it to the Kit camera. + + This mirrors show-renderproduct.py's CopyPrim flow, but uses Sdf.CopySpec so + the RenderProduct can be copied from the source USD package even when it is + not present in the composed GUI stage. + """ + from pxr import Sdf, Usd, UsdRender + + source_rp_prim, source_camera_path = _get_source_render_product(source_stage, source_usd_path) + prim_stack = source_rp_prim.GetPrimStack() + if not prim_stack: + raise RuntimeError(f"PPISP RenderProduct has empty prim stack: {source_rp_prim.GetPath()}") + + source_layer = prim_stack[0].layer + session = stage.GetSessionLayer() + src_path = Sdf.Path(str(source_rp_prim.GetPath())) + dst_path = Sdf.Path(KIT_PPISP_RENDER_PRODUCT_PATH) + package_path = source_stage.GetRootLayer().identifier + + with Usd.EditContext(stage, session): + existing = stage.GetPrimAtPath(dst_path) + if existing and existing.IsValid(): + stage.RemovePrim(dst_path) + + Sdf.CreatePrimInLayer(session, dst_path) + Sdf.CopySpec(source_layer, src_path, session, dst_path) + + with Usd.EditContext(stage, session): + dst_prim = stage.GetPrimAtPath(dst_path) + for prim in Usd.PrimRange(dst_prim): + for attr in prim.GetAttributes(): + connections = attr.GetConnections() + remapped = [conn.ReplacePrefix(src_path, dst_path) for conn in connections] + if remapped != connections: + attr.SetConnections(remapped) + + value = attr.Get() + if isinstance(value, Sdf.AssetPath): + attr.Set(Sdf.AssetPath(_package_qualified_asset(package_path, value.path))) + + for rel in prim.GetRelationships(): + targets = rel.GetTargets() + remapped = [target.ReplacePrefix(src_path, dst_path) for target in targets] + if remapped != targets: + rel.SetTargets(remapped) + + spec = session.GetPrimAtPath(prim.GetPath()) + if spec and spec.referenceList.prependedItems: + spec.referenceList.prependedItems = [ + Sdf.Reference(_package_qualified_asset(package_path, ref.assetPath)) + for ref in spec.referenceList.prependedItems + ] + + UsdRender.Product(dst_prim).GetCameraRel().SetTargets([Sdf.Path(kit_camera_path)]) + + return KIT_PPISP_RENDER_PRODUCT_PATH, source_camera_path + + +def _apply_render_settings_from_stage(stage): + """Apply render settings authored in the source stage custom layer data.""" + import carb.settings + + layer_data = stage.GetRootLayer().customLayerData or {} + render_settings = layer_data.get("renderSettings") or {} + if not render_settings: + return + + settings = carb.settings.get_settings() + for key, value in render_settings.items(): + settings.set("/" + str(key).replace(":", "/"), value) + + +def _copy_camera_metadata_to_chase_camera(stage, source_stage, source_camera_path, kit_camera_path): + """Copy NuRec PPISP camera metadata while preserving the Kit chase-camera pose.""" + from pxr import Sdf, Usd + + source_camera_prim = source_stage.GetPrimAtPath(source_camera_path) + if (not source_camera_prim or not source_camera_prim.IsValid() + or source_camera_prim.GetTypeName() != "Camera"): + raise RuntimeError(f"Viewport source camera does not exist: {source_camera_path}") + + chase_camera_prim = stage.GetPrimAtPath(kit_camera_path) + if (not chase_camera_prim or not chase_camera_prim.IsValid() + or chase_camera_prim.GetTypeName() != "Camera"): + raise RuntimeError(f"Kit chase camera does not exist: {kit_camera_path}") + + with Usd.EditContext(stage, stage.GetSessionLayer()): + api_schemas = source_camera_prim.GetMetadata("apiSchemas") + if api_schemas: + chase_camera_prim.SetMetadata("apiSchemas", api_schemas) + + for src_attr in source_camera_prim.GetAuthoredAttributes(): + name = src_attr.GetName() + if not name.startswith("ppisp:"): + continue + + value = src_attr.Get() + if value is None: + continue + + dst_attr = chase_camera_prim.GetAttribute(name) + if not dst_attr: + dst_attr = chase_camera_prim.CreateAttribute(name, src_attr.GetTypeName()) + dst_attr.Set(value) + + chase_camera_prim.AddAppliedSchema("OmniRtxCameraAutoExposureAPI_1") + chase_camera_prim.AddAppliedSchema("OmniRtxCameraExposureAPI_1") + for name, value in NUREC_IDENTITY_EXPOSURE.items(): + chase_camera_prim.CreateAttribute(name, Sdf.ValueTypeNames.Float).Set(value) + chase_camera_prim.CreateAttribute("omni:rtx:autoExposure:enabled", + Sdf.ValueTypeNames.Bool).Set(False) + + +def _bind_nurec_ppisp_render_product(stage, viewport, nurec_usd_path, kit_camera_path, quiet=False): + """Bind the Kit viewport to a copied NuRec PPISP RenderProduct.""" + if not nurec_usd_path: + return None + + try: + from pxr import Sdf, Usd + + source_stage = Usd.Stage.Open(nurec_usd_path) + if source_stage is None: + raise RuntimeError(f"Could not open NuRec USD: {nurec_usd_path}") + + render_product_path, source_camera_path = _copy_source_ppisp_render_product( + stage, source_stage, nurec_usd_path, kit_camera_path) + _copy_camera_metadata_to_chase_camera(stage, source_stage, source_camera_path, + kit_camera_path) + _apply_render_settings_from_stage(source_stage) + + viewport.camera_path = Sdf.Path(kit_camera_path) + viewport.render_product_path = render_product_path + if str(getattr(viewport, "render_product_path", "")) != render_product_path: + if not quiet: + print(f"[WARN] Kit ignored viewport RenderProduct " + f"{render_product_path}; keeping default viewport render product.") + return None + + return render_product_path + except Exception as exc: # pylint: disable=broad-except + if not quiet: + print(f"[WARN] Could not bind Kit viewport RenderProduct: " + f"{type(exc).__name__}: {exc}") + return None diff --git a/compass/utils/visualizer_utils.py b/compass/utils/visualizer_utils.py new file mode 100644 index 0000000..1b5a457 --- /dev/null +++ b/compass/utils/visualizer_utils.py @@ -0,0 +1,125 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# pylint: disable=import-outside-toplevel + +# Public Exports +__all__ = [ + "KIT_PERSPECTIVE_CAMERA_PATH", + "KIT_SCENE_PARTITION", + "configure_kit_scene_partition", + "configure_visualizers", + "requested_visualizers", +] + +# Kit/Newton Visualizer Defaults +KIT_PERSPECTIVE_CAMERA_PATH = "/OmniverseKit_Persp" +KIT_SCENE_PARTITION = "env_0" +KIT_VIEWER_EYE = (1.5, -1.5, 2.2) +KIT_VIEWER_LOOKAT = (0.0, 0.0, 0.35) + + +# Public API +def requested_visualizers(args): + """Return normalized visualizer names from CLI arguments.""" + requested_viz = getattr(args, "visualizer", None) or [] + if isinstance(requested_viz, str): + requested_viz = requested_viz.split(",") + return [v.strip().lower() for v in requested_viz] + + +def configure_visualizers(env_cfg, requested_viz): + """Configure requested Isaac Lab visualizers.""" + if "newton" in requested_viz or "newton_gl" in requested_viz: + _configure_newton_visualizer(env_cfg) + if "kit" in requested_viz: + _configure_kit_visualizer(env_cfg) + + +# Private Helpers +def _append_visualizer_cfg(env_cfg, visualizer_cfg): + """Append a visualizer config while preserving existing config shape.""" + existing_cfgs = env_cfg.sim.visualizer_cfgs + if existing_cfgs is None: + env_cfg.sim.visualizer_cfgs = [visualizer_cfg] + elif isinstance(existing_cfgs, list): + env_cfg.sim.visualizer_cfgs.append(visualizer_cfg) + else: + env_cfg.sim.visualizer_cfgs = [existing_cfgs, visualizer_cfg] + + +################################ +# Kit Visualizer Setup +################################ +def configure_kit_scene_partition(quiet=False): + """Restrict the Kit perspective camera to the first environment clone.""" + try: + import omni.usd # pylint: disable=import-outside-toplevel + from pxr import Sdf # pylint: disable=import-outside-toplevel + except ImportError as exc: + if not quiet: + print(f"[WARN] Could not import USD utilities for Kit viewport setup: {exc}") + return + + stage = omni.usd.get_context().get_stage() + if stage is None: + if not quiet: + print("[WARN] Could not configure Kit viewport: no active USD stage.") + return + + camera_prim = stage.GetPrimAtPath(KIT_PERSPECTIVE_CAMERA_PATH) + if not camera_prim or not camera_prim.IsValid(): + if not quiet: + print(f"[WARN] Could not configure Kit viewport: " + f"{KIT_PERSPECTIVE_CAMERA_PATH} does not exist.") + return + + attr = camera_prim.GetAttribute("omni:scenePartition") + if not attr.IsValid(): + attr = camera_prim.CreateAttribute("omni:scenePartition", Sdf.ValueTypeNames.Token) + attr.Set(KIT_SCENE_PARTITION) + if not quiet: + print(f"[INFO] Set {KIT_PERSPECTIVE_CAMERA_PATH} scene partition to " + f"{KIT_SCENE_PARTITION}.") + + +def _configure_kit_visualizer(env_cfg): + """Configure the Kit visualizer follow camera.""" + from isaaclab_visualizers.kit import KitVisualizerCfg + + kit_viz_cfg = KitVisualizerCfg() + kit_viz_cfg.eye = KIT_VIEWER_EYE + kit_viz_cfg.lookat = KIT_VIEWER_LOOKAT + kit_viz_cfg.origin_type = "asset" + kit_viz_cfg.origin_track_path = "robot" + kit_viz_cfg.origin_env_index = 0 + _append_visualizer_cfg(env_cfg, kit_viz_cfg) + + +################################ +# Newton Visualizer Setup +################################ +def _configure_newton_visualizer(env_cfg): + """Configure Newton visualizer framing and follow-camera panel.""" + from isaaclab_visualizers.newton import NewtonVisualizerCfg + + newton_viz_cfg = NewtonVisualizerCfg() + newton_viz_cfg.eye = KIT_VIEWER_EYE + newton_viz_cfg.lookat = KIT_VIEWER_LOOKAT + newton_viz_cfg.tiled_cam_view = True + newton_viz_cfg.tiled_cam_num = 1 + newton_viz_cfg.tiled_cam_target_prim_path = "/World/envs/*/Robot" + newton_viz_cfg.tiled_cam_eye = KIT_VIEWER_EYE + _append_visualizer_cfg(env_cfg, newton_viz_cfg) diff --git a/docker/prepare_assets.sh b/docker/prepare_assets.sh index 95eb6b7..0d02625 100755 --- a/docker/prepare_assets.sh +++ b/docker/prepare_assets.sh @@ -16,6 +16,7 @@ # Environment: # HF_TOKEN HuggingFace token (https://huggingface.co/settings/tokens) # Required because the COMPASS HF repo is gated. The flag overrides. +# NUREC_REVISION Optional NuRec dataset revision. Empty uses the HF default branch. set -eu @@ -36,7 +37,7 @@ CACHE_DIR="${REPO_ROOT}/assets" HF_TOKEN_ARG="" FORCE=false NUREC_SCENES=() -NUREC_REVISION="${NUREC_REVISION:-refs/pr/34}" +NUREC_REVISION="${NUREC_REVISION:-}" USDS_URL="https://huggingface.co/nvidia/COMPASS/resolve/main/compass_usds.zip" USDS_ZIP="${CACHE_DIR}/compass_usds.zip" @@ -220,10 +221,17 @@ download_nurec_scenes() { continue fi - step "Downloading NuRec scene ${scene} from ${NUREC_REPO}@${NUREC_REVISION} → ${USDS_DIR}" + local revision_args=() + local revision_label="dataset default revision" + if [ -n "${NUREC_REVISION}" ]; then + revision_args=(--revision "${NUREC_REVISION}") + revision_label="${NUREC_REPO}@${NUREC_REVISION}" + fi + + step "Downloading NuRec scene ${scene} from ${revision_label} → ${USDS_DIR}" HF_TOKEN="${HF_TOKEN_EFFECTIVE}" hf download "${NUREC_REPO}" \ --repo-type dataset \ - --revision "${NUREC_REVISION}" \ + "${revision_args[@]}" \ --local-dir "${USDS_DIR}" \ --include "${scene}/**" \ --exclude "**/raw_images.zip" @@ -245,7 +253,11 @@ show_summary() { echo " ✗ X-Mobility ckpt: ${CKPT_FILE} (missing)" fi if [ "${#NUREC_SCENES[@]}" -gt 0 ]; then - echo " NuRec revision: ${NUREC_REPO}@${NUREC_REVISION}" + if [ -n "${NUREC_REVISION}" ]; then + echo " NuRec revision: ${NUREC_REPO}@${NUREC_REVISION}" + else + echo " NuRec revision: ${NUREC_REPO} default branch" + fi local scene for scene in "${NUREC_SCENES[@]}"; do echo " ✓ NuRec scene: ${USDS_DIR}/${scene}" diff --git a/docs/handbook/workflows/images/nurec_sh_optimized_camera_grid.png b/docs/handbook/workflows/images/nurec_sh_optimized_camera_grid.png deleted file mode 100644 index 52c3daa..0000000 --- a/docs/handbook/workflows/images/nurec_sh_optimized_camera_grid.png +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:d0a5731ae32f358c4a42a90e5a23d50daa8216b7839814a48553b9e6af609ce3 -size 554265 diff --git a/docs/handbook/workflows/images/nurec_sh_optimized_kit_viewport.jpg b/docs/handbook/workflows/images/nurec_sh_optimized_kit_viewport.jpg deleted file mode 100644 index cf3a3e3..0000000 Binary files a/docs/handbook/workflows/images/nurec_sh_optimized_kit_viewport.jpg and /dev/null differ diff --git a/docs/handbook/workflows/nurec_real2sim.md b/docs/handbook/workflows/nurec_real2sim.md index 16f594b..6e11792 100644 --- a/docs/handbook/workflows/nurec_real2sim.md +++ b/docs/handbook/workflows/nurec_real2sim.md @@ -45,12 +45,12 @@ installs requested NuRec scenes under export HF_TOKEN=hf_xxx ./docker/run.sh assets \ - --nurec-scene nova_carter-galileo \ - --nurec-revision refs/pr/34 + --nurec-scene nova_carter-galileo ``` -Use one `--nurec-scene ` flag per scene. The helper excludes -`raw_images.zip`, skips existing assets, and is safe to re-run. +Use one `--nurec-scene ` flag per scene. +Pass `--nurec-revision ` to download Hugging Face data from a +specific branch, tag, commit, or PR ref. ```{note} You must accept the dataset terms on Hugging Face before downloading. @@ -58,10 +58,6 @@ You must accept the dataset terms on Hugging Face before downloading. ### 4. Build and activate the container -The NuRec image currently uses an internal Isaac Lab 3.0 release base image -because the old public `3.0.0-beta1` image is not compatible with this branch. -Replace it with the matching public Isaac Lab image after release. - ```bash export COMPASS_IMAGE_TAG=isaaclab-3.0-ea @@ -81,33 +77,22 @@ Each NuRec scene should live at: compass/rl_env/exts/mobility_es/mobility_es/usd// ``` -The asset helper installs scenes there automatically. A scene folder is flat: +The asset helper installs scenes there automatically. Relevant scene files +include: ```text nova_carter-galileo/ -├── particle_spg-runtime.usdz -├── particle_sh_optimized.usdz +├── particle_spg-runtime.usdz # Particle-field USD with SPG ├── volume.usdz ├── occupancy_map.yaml ├── occupancy_map.png ``` -Scene names are listed in `compass/utils/nurec_utils.py`. - -Registered scenes: - -| `--nurec-scene` | Description | -|---|---| -| `nova_carter-galileo` | Galileo lab | -| `nova_carter-cafe` | NVIDIA cafe | -| `hand_hold-endeavor-andoria` | Endeavor meeting room | -| `hand_hold-endeavor-livingroom` | Endeavor living room | -| `hand_hold-endeavor-wormhole` | Endeavor conference room | -| `hand_hold-endeavor-wormhole-table` | Endeavor conference room with table | -| `hand_hold-voyager-babyboom` | Voyager conference room | +Pass the scene folder name to `--nurec-scene`; COMPASS builds the NuRec scene +config from that folder at runtime. ```{note} -NuRec maps use a ROS bottom-left origin convention. The registry sets this by +NuRec maps use a ROS bottom-left origin convention. COMPASS sets this by default; using the wrong convention leads to bad robot spawn locations. ``` @@ -121,15 +106,6 @@ cd - ### 7. Train -Use `--nurec-usd-file ` to choose a USD from the selected scene -folder. The current workflow supports particle USD assets; volume -assets are present in some scene folders but are not tested. - -| Asset | Use | Notes | -|---|---|---| -| `particle_spg-runtime.usdz` | Default particle USD. | Automatically turns on SPG in Isaac Sim and enables the copied PPISP viewport path. | -| `particle_sh_optimized.usdz` | Particle USD without SPG runtime. | SPG runtime stays off automatically; no extra flag is needed. | - #### Default particle SPG-runtime command: ```bash @@ -142,38 +118,16 @@ python run.py \ --num_envs 12 \ --video \ --video_interval 1 \ - --visualizer kit \ - --precompute_valid_poses -``` - -Reference output: - -| Kit viewport | Robot-camera debug grid | -|---|---| -| Particle SPG-runtime Kit viewport | Particle SPG-runtime robot-camera debug grid | - -#### SH-optimized particle command: - -```bash -python run.py \ - -c configs/train_config_real2sim.gin \ - -o \ - -b ./assets/x_mobility.ckpt \ - --embodiment carter \ - --nurec-scene nova_carter-galileo \ - --nurec-usd-file particle_sh_optimized.usdz \ - --num_envs 12 \ - --video \ - --video_interval 1 \ - --visualizer kit \ --precompute_valid_poses ``` -Reference output: +This command uses the default `particle_spg-runtime.usdz`. Use +`--nurec-usd-file ` to choose another USD from the selected scene +folder. -| Kit viewport | Robot-camera debug grid | +| Kit viewport | Robot camera | |---|---| -| SH-optimized Kit viewport | SH-optimized robot-camera debug grid | +| Particle SPG-runtime Kit viewport | Particle SPG-runtime robot-camera debug grid | The Kit viewport is an extreme novel view from above the robot, so render quality can be lower than the onboard camera. It is only for user visualization @@ -183,38 +137,35 @@ policy-observation reference used by training. Key options: - ``: `h1`, `spot`, `carter`, `g1`, or `digit`. -- `--nurec-scene`: any registered NuRec scene. This is a NuRec-specific alias - for `--environment`; assets must already be installed. +- `--nurec-scene`: any installed NuRec scene folder. This is a NuRec-specific + alias for `--environment`; assets must already be installed. - `--nurec-usd-file`: NuRec USD filename under the selected scene folder. - Defaults to `particle_spg-runtime.usdz`; use - `--nurec-usd-file particle_sh_optimized.usdz` to swap assets. + Defaults to `particle_spg-runtime.usdz`; override it only when testing a + custom NuRec USD. +- `--nurec-omap-file`: optional occupancy-map YAML filename under the selected + scene folder. Omit it to use the scene default, usually + `occupancy_map.yaml`. - `--spg-runtime`: advanced override for custom SPG-runtime USDs. Usually omit - it; COMPASS automatically enables SPG for `particle_spg-runtime.usdz` in - registered NuRec scenes and leaves it off for `particle_sh_optimized.usdz`. + it; COMPASS automatically enables SPG for `particle_spg-runtime.usdz` when + `--nurec-scene` is set. - `--num_envs 12`: conservative default; increase only after checking VRAM. - `--precompute_valid_poses`: recommended for constrained Real2Sim scenes. -- With `--visualizer kit`, the GUI uses the Kit perspective camera; debug - images and policy observations still use the robot camera sensor. For - `particle_spg-runtime.usdz`, COMPASS copies the authored NuRec PPISP - RenderProduct, retargets it to the chase camera, and applies NuRec identity - exposure. SPG runtime Kit args are enabled automatically for this asset - before Isaac Sim starts. COMPASS uses the first RenderProduct with a camera - target in the source NuRec USD. -- Debug dumps include `camera_grid_*.png` for robot-camera RGB/depth and - `kit_viewport_*.png` plus `kit_viewport_*.png.txt` for the GUI viewport. +- Debug dumps include `camera_grid_*.png` for robot-camera RGB/depth. With + `--visualizer kit`, they also include `kit_viewport_*.png` plus + `kit_viewport_*.png.txt` for the GUI viewport. +- Use `--visualizer kit` to show third-person visualization, but it can be much + slower. Omit it for training. Output checkpoints are written as `/model_.pt`; videos land in `/videos/`. ```{note} -To run headless, omit `--visualizer kit` or pass `--visualizer None`. -``` +Approximate per-GPU VRAM for `nova_carter-galileo` with Carter and a 320x512 +camera: -```{note} -Approximate per-GPU VRAM for `nova_carter-galileo`, Carter, 320x512 camera: -particle assets use `8.5 GB + 1.0 GB × num_envs`. -Safe particle defaults are about 18 envs on 32 GB and 32 envs on 48 GB. -Reduce `--num_envs` on OOM. +- Particle assets: `8.5 GB + 1.0 GB * num_envs` +- Safe defaults: about 18 envs on 32 GB, or 32 envs on 48 GB +- OOM recovery: reduce `--num_envs` ``` ### 8. Train on multiple GPUs @@ -248,8 +199,7 @@ python run.py \ --nurec-scene nova_carter-galileo \ --num_envs \ --video \ - --video_interval 1 \ - --visualizer kit + --video_interval 1 ``` `` is usually `/model_.pt`. @@ -288,19 +238,19 @@ osmo login NuRec OSMO jobs download COMPASS USDs, the X-Mobility checkpoint, and the requested NuRec scene inside the workflow. When `--nurec-scene` is set, the workflow switches to the Real2Sim gin config and passes `--nurec-scene` and -`--nurec-usd-file` to `run.py`. -Use `--nurec-scene` instead of `--environment` so the same scene drives both -the Hugging Face asset download and the NuRec runtime branch. +`--nurec-usd-file` to `run.py`. Use `--nurec-omap-file ` when the +workflow should use a non-default occupancy map. Omit `--nurec-revision` to use +the Hugging Face dataset default branch; pass it only to pin a specific +revision. Training example: ```bash python osmo/run_osmo.py train \ - --experiment-name nurec \ - --wandb-project compass-nurec-carter-galileo \ + --experiment-name carter-galileo-osmo \ + --wandb-project compass-nurec \ --embodiment carter \ --nurec-scene nova_carter-galileo \ - --nurec-revision refs/pr/34 \ --nurec-usd-file particle_spg-runtime.usdz \ --num-envs 12 \ --num-gpus 8 @@ -310,12 +260,11 @@ Evaluation example: ```bash python osmo/run_osmo.py eval \ - --experiment-name nurec-eval \ - --wandb-project compass-nurec-carter-galileo-eval \ + --experiment-name carter-galileo-osmo-eval \ + --wandb-project compass-nurec \ --checkpoint \ --embodiment carter \ --nurec-scene nova_carter-galileo \ - --nurec-revision refs/pr/34 \ --nurec-usd-file particle_spg-runtime.usdz \ --num-envs 12 ``` diff --git a/osmo/run_osmo.py b/osmo/run_osmo.py index c90606b..9ac9717 100755 --- a/osmo/run_osmo.py +++ b/osmo/run_osmo.py @@ -64,7 +64,6 @@ REPO_ROOT = Path(__file__).resolve().parent.parent WORKFLOWS_DIR = Path(__file__).resolve().parent / "workflows" -DEFAULT_NUREC_REVISION = "refs/pr/34" # subcommand -> (workflow YAML filename, Dockerfile path relative to REPO_ROOT) SUBCOMMAND_CONFIG = { @@ -113,11 +112,16 @@ def add_nurec(sp): "For NuRec runs this also selects the runtime environment.") sp.add_argument( "--nurec-revision", - default=DEFAULT_NUREC_REVISION, - help="Git revision in nvidia/PhysicalAI-Robotics-NuRec used for NuRec assets.") + default="", + help="Optional Git revision in nvidia/PhysicalAI-Robotics-NuRec used for NuRec " + "assets. Empty uses the dataset repo default branch.") sp.add_argument("--nurec-usd-file", default="", help="NuRec USD filename passed to run.py. Empty uses run.py's default.") + sp.add_argument("--nurec-omap-file", + default="", + help="NuRec occupancy-map filename passed to run.py. " + "Empty uses run.py's default.") train = sub.add_parser("train", help="Submit residual RL training.") add_common(train) @@ -264,6 +268,7 @@ def cmd_train(args, image: str, wandb_key: str, hf_token: str) -> None: "nurec_scene": args.nurec_scene, "nurec_revision": args.nurec_revision, "nurec_usd_file": args.nurec_usd_file, + "nurec_omap_file": args.nurec_omap_file, } submit_workflow(yaml_path, set_args, args.dry_run) @@ -287,6 +292,7 @@ def cmd_eval(args, image: str, wandb_key: str, hf_token: str) -> None: "nurec_scene": args.nurec_scene, "nurec_revision": args.nurec_revision, "nurec_usd_file": args.nurec_usd_file, + "nurec_omap_file": args.nurec_omap_file, } submit_workflow(yaml_path, set_args, args.dry_run) diff --git a/osmo/workflows/rl_es_eval_workflow.yaml b/osmo/workflows/rl_es_eval_workflow.yaml index 9c7fbbd..6f24c0f 100644 --- a/osmo/workflows/rl_es_eval_workflow.yaml +++ b/osmo/workflows/rl_es_eval_workflow.yaml @@ -51,6 +51,7 @@ workflow: NUREC_SCENE="{{nurec_scene}}" NUREC_REVISION="{{nurec_revision}}" NUREC_USD_FILE="{{nurec_usd_file}}" + NUREC_OMAP_FILE="{{nurec_omap_file}}" EVAL_CONFIG=configs/eval_config.gin # Download and unpack the COMPASS USDs from HuggingFace. @@ -65,10 +66,16 @@ workflow: if [[ -z "$ENVIRONMENT" ]]; then ENVIRONMENT="$NUREC_SCENE" fi - echo "Downloading $NUREC_SCENE NuRec assets from $NUREC_REVISION" + NUREC_REVISION_ARGS=() + if [[ -n "$NUREC_REVISION" ]]; then + echo "Downloading $NUREC_SCENE NuRec assets from $NUREC_REVISION" + NUREC_REVISION_ARGS=(--revision "$NUREC_REVISION") + else + echo "Downloading $NUREC_SCENE NuRec assets from the dataset default revision" + fi ${ISAACLAB_PATH}/isaaclab.sh -p -m huggingface_hub.commands.huggingface_cli \ download nvidia/PhysicalAI-Robotics-NuRec --repo-type dataset \ - --revision "$NUREC_REVISION" \ + "${NUREC_REVISION_ARGS[@]}" \ --local-dir ./compass/rl_env/exts/mobility_es/mobility_es/usd \ --include "$NUREC_SCENE/**" \ --exclude "**/raw_images.zip" @@ -114,6 +121,9 @@ workflow: if [[ -n $NUREC_USD_FILE ]]; then CMD+=" --nurec-usd-file $NUREC_USD_FILE" fi + if [[ -n $NUREC_OMAP_FILE ]]; then + CMD+=" --nurec-omap-file $NUREC_OMAP_FILE" + fi if [[ -n $DESTILLATION_CKPT_ARTIFACT ]]; then CMD+=" -d $DESTILLATION_CKPT" @@ -144,5 +154,6 @@ default-values: embodiment: "" environment: "" nurec_scene: "" - nurec_revision: refs/pr/34 + nurec_revision: "" nurec_usd_file: "" + nurec_omap_file: "" diff --git a/osmo/workflows/rl_es_train_workflow.yaml b/osmo/workflows/rl_es_train_workflow.yaml index 5b8b666..951e288 100644 --- a/osmo/workflows/rl_es_train_workflow.yaml +++ b/osmo/workflows/rl_es_train_workflow.yaml @@ -64,6 +64,7 @@ workflow: NUREC_SCENE="{{nurec_scene}}" NUREC_REVISION="{{nurec_revision}}" NUREC_USD_FILE="{{nurec_usd_file}}" + NUREC_OMAP_FILE="{{nurec_omap_file}}" TRAIN_CONFIG=configs/train_config.gin EVAL_CONFIG=configs/eval_config.gin @@ -79,10 +80,16 @@ workflow: if [[ -z "$ENVIRONMENT" ]]; then ENVIRONMENT="$NUREC_SCENE" fi - echo "Downloading $NUREC_SCENE NuRec assets from $NUREC_REVISION" + NUREC_REVISION_ARGS=() + if [[ -n "$NUREC_REVISION" ]]; then + echo "Downloading $NUREC_SCENE NuRec assets from $NUREC_REVISION" + NUREC_REVISION_ARGS=(--revision "$NUREC_REVISION") + else + echo "Downloading $NUREC_SCENE NuRec assets from the dataset default revision" + fi python -m huggingface_hub.commands.huggingface_cli \ download nvidia/PhysicalAI-Robotics-NuRec --repo-type dataset \ - --revision "$NUREC_REVISION" \ + "${NUREC_REVISION_ARGS[@]}" \ --local-dir ./compass/rl_env/exts/mobility_es/mobility_es/usd \ --include "$NUREC_SCENE/**" \ --exclude "**/raw_images.zip" @@ -127,6 +134,9 @@ workflow: if [[ -n $NUREC_USD_FILE ]]; then TRAIN_CMD+=" --nurec-usd-file $NUREC_USD_FILE" fi + if [[ -n $NUREC_OMAP_FILE ]]; then + TRAIN_CMD+=" --nurec-omap-file $NUREC_OMAP_FILE" + fi if [[ -z $NO_RESIDUAL ]]; then TRAIN_CMD+=" -b ./model.ckpt" @@ -162,6 +172,9 @@ workflow: if [[ -n $NUREC_USD_FILE ]]; then EVAL_CMD+=" --nurec-usd-file $NUREC_USD_FILE" fi + if [[ -n $NUREC_OMAP_FILE ]]; then + EVAL_CMD+=" --nurec-omap-file $NUREC_OMAP_FILE" + fi echo $EVAL_CMD eval $EVAL_CMD @@ -185,5 +198,6 @@ default-values: embodiment: "" environment: "" nurec_scene: "" - nurec_revision: refs/pr/34 + nurec_revision: "" nurec_usd_file: "" + nurec_omap_file: "" diff --git a/run.py b/run.py index 894008e..9117ecc 100644 --- a/run.py +++ b/run.py @@ -21,79 +21,17 @@ from isaaclab.app import AppLauncher -from compass.utils.nurec_utils import ( - PARTICLE_SPG_RUNTIME_USD_FILE, +from compass.utils.nurec_utils import PARTICLE_SPG_RUNTIME_USD_FILE +from compass.utils.renderer_utils import ( + apply_isaac_rtx_camera_renderer_settings, apply_nurec_spg_kit_args, - configure_nurec_kit_viewport, + configure_nurec_isaacsim_rtx_viewport, +) +from compass.utils.visualizer_utils import ( + configure_kit_scene_partition, + configure_visualizers, + requested_visualizers, ) - -KIT_PERSPECTIVE_CAMERA_PATH = "/OmniverseKit_Persp" -KIT_SCENE_PARTITION = "env_0" -KIT_VIEWER_EYE = (1.5, -1.5, 2.2) -KIT_VIEWER_LOOKAT = (0.0, 0.0, 0.35) - - -def _append_visualizer_cfg(env_cfg, visualizer_cfg): - existing_cfgs = env_cfg.sim.visualizer_cfgs - if existing_cfgs is None: - env_cfg.sim.visualizer_cfgs = [visualizer_cfg] - elif isinstance(existing_cfgs, list): - env_cfg.sim.visualizer_cfgs.append(visualizer_cfg) - else: - env_cfg.sim.visualizer_cfgs = [existing_cfgs, visualizer_cfg] - - -def _configure_legacy_viewer(env_cfg): - env_cfg.viewer.origin_type = "asset_root" - env_cfg.viewer.asset_name = "robot" - env_cfg.viewer.env_index = 0 - env_cfg.viewer.cam_prim_path = KIT_PERSPECTIVE_CAMERA_PATH - env_cfg.viewer.eye = KIT_VIEWER_EYE - env_cfg.viewer.lookat = KIT_VIEWER_LOOKAT - - -def _configure_kit_visualizer(env_cfg): - from isaaclab_visualizers.kit import KitVisualizerCfg - - kit_viz_cfg = KitVisualizerCfg() - kit_viz_cfg.eye = KIT_VIEWER_EYE - kit_viz_cfg.lookat = KIT_VIEWER_LOOKAT - kit_viz_cfg.origin_type = "asset" - kit_viz_cfg.origin_track_path = "robot" - kit_viz_cfg.origin_env_index = 0 - _append_visualizer_cfg(env_cfg, kit_viz_cfg) - - -def _configure_kit_scene_partition(quiet=False): - try: - import omni.usd - from pxr import Sdf - except ImportError as exc: - if not quiet: - print(f"[WARN] Could not import USD utilities for Kit viewport setup: {exc}") - return - - stage = omni.usd.get_context().get_stage() - if stage is None: - if not quiet: - print("[WARN] Could not configure Kit viewport: no active USD stage.") - return - - camera_prim = stage.GetPrimAtPath(KIT_PERSPECTIVE_CAMERA_PATH) - if not camera_prim or not camera_prim.IsValid(): - if not quiet: - print(f"[WARN] Could not configure Kit viewport: " - f"{KIT_PERSPECTIVE_CAMERA_PATH} does not exist.") - return - - attr = camera_prim.GetAttribute("omni:scenePartition") - if not attr.IsValid(): - attr = camera_prim.CreateAttribute("omni:scenePartition", Sdf.ValueTypeNames.Token) - attr.Set(KIT_SCENE_PARTITION) - if not quiet: - print(f"[INFO] Set {KIT_PERSPECTIVE_CAMERA_PATH} scene partition to " - f"{KIT_SCENE_PARTITION}.") - # add argparse arguments parser = argparse.ArgumentParser(description="COMPASS Mobility Generalist.") @@ -190,21 +128,23 @@ def _configure_kit_scene_partition(quiet=False): ) parser.add_argument( "--nurec-usd-file", - "--nurec_usd_file", - dest="nurec_usd_file", type=str, default=PARTICLE_SPG_RUNTIME_USD_FILE, help="NuRec USD filename under the selected environment folder.", ) +parser.add_argument( + "--nurec-omap-file", + type=str, + default=None, + help="NuRec occupancy-map YAML filename under the selected environment folder.", +) parser.add_argument( "--spg-runtime", - "--spg_runtime", - dest="spg_runtime", action="store_true", default=False, help="Force SPG runtime Kit args before Isaac Sim starts. " f"This is automatic for {PARTICLE_SPG_RUNTIME_USD_FILE} " - "in registered NuRec scenes.", + "when --nurec-scene is set.", ) parser.add_argument("--num_envs", type=int, help="Number of environments") parser.add_argument( @@ -245,6 +185,8 @@ def _configure_kit_scene_partition(quiet=False): if args_cli.environment is not None: parser.error("Pass either --nurec-scene or --environment, not both.") args_cli.environment = args_cli.nurec_scene +elif args_cli.nurec_omap_file is not None: + parser.error("--nurec-omap-file requires --nurec-scene.") apply_nurec_spg_kit_args(args_cli) # launch omniverse app @@ -262,7 +204,7 @@ def _configure_kit_scene_partition(quiet=False): from mobility_es.config.spot_env_cfg import SpotGoalReachingEnvCfg from mobility_es.config.g1_env_cfg import G1GoalReachingEnvCfg from mobility_es.config.digit_env_cfg import DigitGoalReachingEnvCfg -from mobility_es.config.nurec_scenes import make_nurec_scene_asset_cfg_map +from mobility_es.config.nurec_scenes import make_nurec_scene_asset_cfg from mobility_es.wrapper.env_wrapper import RLESEnvWrapper from compass.residual_rl.x_mobility_rl import XMobilityBasePolicy @@ -312,15 +254,12 @@ def close(self): "hospital": environments.hospital, "warehouse_multi_rack": environments.warehouse_multi_rack, } -NurecSceneAssetCfgMap = make_nurec_scene_asset_cfg_map(args_cli.nurec_usd_file) -EnvSceneAssetCfgMap.update(NurecSceneAssetCfgMap) - - -def _requested_visualizers(): - requested_viz = getattr(args_cli, "visualizer", None) or [] - if isinstance(requested_viz, str): - requested_viz = requested_viz.split(",") - return [v.strip().lower() for v in requested_viz] +if args_cli.nurec_scene is not None: + EnvSceneAssetCfgMap[args_cli.nurec_scene] = make_nurec_scene_asset_cfg( + args_cli.nurec_scene, + args_cli.nurec_usd_file, + args_cli.nurec_omap_file, + ) def gin_config_to_dictionary(gin_config): @@ -448,25 +387,13 @@ def run( else: env_cfg.curriculum = None - _configure_legacy_viewer(env_cfg) - - # Newton visualizer camera (new Visualizers API). Newton reads its camera from - # sim.visualizer_cfgs, not ViewerCfg. Kit also uses visualizer_cfgs in Isaac Lab 3, - # which keeps the GUI perspective camera independent from the onboard robot camera. - _requested_viz = _requested_visualizers() - if "newton" in _requested_viz or "newton_gl" in _requested_viz: - from isaaclab_visualizers.newton import NewtonVisualizerCfg - - newton_viz_cfg = NewtonVisualizerCfg() - newton_viz_cfg.eye = KIT_VIEWER_EYE # initial interactive framing - newton_viz_cfg.lookat = KIT_VIEWER_LOOKAT - newton_viz_cfg.tiled_cam_view = True # follow-cam panel (Newton's follow path) - newton_viz_cfg.tiled_cam_num = 1 - newton_viz_cfg.tiled_cam_target_prim_path = "/World/envs/*/Robot" - newton_viz_cfg.tiled_cam_eye = KIT_VIEWER_EYE - _append_visualizer_cfg(env_cfg, newton_viz_cfg) - if "kit" in _requested_viz: - _configure_kit_visualizer(env_cfg) + requested_viz = requested_visualizers(args_cli) + configure_visualizers(env_cfg, requested_viz) + apply_isaac_rtx_camera_renderer_settings( + env_cfg, + isaac_rtx=True, + spg_runtime=args_cli.spg_runtime, + ) # Setup seed. Per-rank offset diversifies env initial conditions across GPUs so # rollouts collected by each rank explore different states (matches Isaac Lab's @@ -505,14 +432,14 @@ def run( precompute_valid_poses=precompute_flag, precompute_valid_orientations=precompute_orientations_flag, ) - is_nurec_run = args_cli.nurec_scene is not None - if not is_nurec_run and "kit" in _requested_viz: - _configure_kit_scene_partition() - if is_nurec_run and "kit" in _requested_viz: - configure_nurec_kit_viewport( - nurec_usd_path=env_cfg.scene.environment.spawn.usd_path, - spg_runtime=args_cli.spg_runtime, - ) + if "kit" in requested_viz: + if args_cli.nurec_scene is not None: + configure_nurec_isaacsim_rtx_viewport( + nurec_usd_path=env_cfg.scene.environment.spawn.usd_path, + spg_runtime=args_cli.spg_runtime, + ) + else: + configure_kit_scene_partition() # Precompute valid pose locations if requested if precompute_flag and env.collision_checker.is_initialized(): @@ -549,6 +476,7 @@ def run( output_dir=args_cli.output_dir, logger=logger, device=device, + save_debug_viewport_images="kit" in requested_viz, ) if run_mode == "train":