scene.py

#
#  Copyright 2024 Proximity Robotics & Automation GmbH
#  
#  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.
#  

from isaacsim import SimulationApp
import scripts.utils
from typing import Optional


class Scene:

    def __init__(
        self,
        display_options: Optional[int] = None,
        data_root: str = "/isaac-sim/project/data",
        scripts_root: str = "/isaac-sim/scripts",
        omniverse_assets_root: str = "http://omniverse-content-production.s3-us-west-2.amazonaws.com",
        add_ground_plane: bool = True,
        add_default_light: bool = False,
        config_filename: str = "simulation_app_config.json",
    ) -> None:
        """Starts the simulation by creating a SimulationApp object and creates a scene in a new stage in Isaac Sim.

        :param display_options: Display options code to specify what's visible in the stage by default. The desired code
         can be calculated with the "get_display_options_code" function in the utils.py script. If None, the
         "get_display_options_code" default is used. Defaults to None.
        :type display_options: Optional[int], optional
        :param data_root: Path to the data folder containing assets etc. in the docker container. Defaults to
         "/isaac-sim/project/data".
        :type data_root: str, optional
        :param scripts_root: Path to the scripts folder in der the docker container. Defaults to "/isaac-sim/scripts".
        :type scripts_root: str, optional
        :param omniverse_assets_root: Path to the Isaac Sim assets provided by NVIDIA. Defaults to
         "http://omniverse-content-production.s3-us-west-2.amazonaws.com".
        :type omniverse_assets_root: str, optional
        :param add_ground_plane: Whether to add a ground plane to the scene. Defaults to True.
        :type add_ground_plane: bool, optional
        :param add_default_light: Whether to add a default DomeLight to the scene. Defaults to False.
        :type add_default_light: bool, optional
        :param config_filename: Filename of the config file to use. The file needs to be stored in "data/config_files/".
         It contains parameters to initialize the SimulationApp object. Defaults to "simulation_app_config.json".
        :type config_filename: str, optional
        """
        self.omniverse_assets_root = omniverse_assets_root
        self.data_root = data_root
        self.scripts_root = scripts_root
        simulation_config = scripts.utils.dict_from_json(
            self.data_root + "/config_files/" + config_filename
        )
        self.simulation = SimulationApp(simulation_config)

        # Manually set "displayOptions" setting as it is somehow not used when given as arg in SimulationApp init
        # TODO: Parameter "hide_ui" as well? Or remove from below?
        # TODO: Favor config file or display_options parameter? / remove display_options parameter and only use config?
        if display_options is None:
            display_options = (
                simulation_config.get("display_options")
                or scripts.utils.get_display_options_code()
            )

        self.simulation.set_setting(
            "persistent/app/viewport/displayOptions", display_options
        )

        if "hide_ui" in simulation_config:
            self.simulation.set_setting(
                "/app/window/hideUi", simulation_config["hide_ui"]
            )

        # Default livestream settings
        self.simulation.set_setting("/app/window/drawMouse", False)
        self.simulation.set_setting("/app/livestream/proto", "ws")
        self.simulation.set_setting("/app/livestream/websocket/framerate_limit", 120)
        self.simulation.set_setting("/ngx/enabled", False)
        # self.simulation.set_setting("/app/window/hideUi", False)

        # Imports
        global isaac_core, carb, transform_utils
        global physx_utils, get_context, UsdGeom, UsdPhysics, Usd, new_stage_with_callback, UsdShade

        from omni.usd import get_context
        import omni.isaac.core as isaac_core
        import carb
        from scripts.isaac import transform_utils

        from omni.physx.scripts import utils as physx_utils
        from pxr import UsdGeom, UsdPhysics, Usd, UsdShade
        import omni.kit.app
        from omni.kit.stage_template.core import new_stage_with_callback
        import omni.kit.window.file

        # Enable "RTX - Scientific (IndeX)" renderer
        # isaac_core.utils.extensions.enable_extension("omni.hydra.index")
        # isaac_core.utils.extensions.enable_extension("omni.index.usd")
        # carb.settings.get_settings().set_string("/renderer/enabled", "rtx,index")

        # Enable "RTX - Accurate (Iray)" renderer
        # isaac_core.utils.extensions.enable_extension("omni.hydra.iray")
        # carb.settings.get_settings().set_string("/renderer/enabled", "rtx,iray")

        # Use WebRTC streaming (http://localhost:8211/streaming/webrtc-client/)
        # isaac_core.utils.extensions.enable_extension("omni.services.streamclient.webrtc")

        extension_manager = omni.kit.app.get_app().get_extension_manager()
        carb_settings = carb.settings.get_settings()

        # Use WebRTC streaming
        extension_manager.set_extension_enabled(
            "omni.services.streamclient.webrtc", True
        )

        self.world = None

        # Enable wheeled robots extension
        isaac_core.utils.extensions.enable_extension("omni.isaac.wheeled_robots")

        # Add LiDAR config path
        lidar_config_folders = carb_settings.get(
            "/app/sensors/nv/lidar/profileBaseFolder"
        )
        lidar_config_folders.append(self.data_root + "/config_files/lidar_configs/")
        carb_settings.set(
            "/app/sensors/nv/lidar/profileBaseFolder", lidar_config_folders
        )

        # Get the current stage and create the world
        self.stage = get_context().get_stage()
        self.world = self.create_world(add_ground_plane, add_default_light)

    def new_stage(self):
        """Opens a new stage.

        :return: The new stage prim.
        :rtype: Usd.Stage
        """
        new_stage_with_callback(None, "empty", get_context())
        self.stage = get_context().get_stage()

    def create_world(self, add_ground_plane: bool, add_default_light: bool):
        """Creates the world prim and potentially adds a ground plane and default light.

        :param add_ground_plane: Whether to add a ground plane.
        :type add_ground_plane: bool
        :param add_default_light: Whether to add a default light (DomeLight).
        :type add_default_light: bool
        :return: The world prim.
        :rtype: omni.isaac.core.World
        """
        # Create the world
        world = isaac_core.World(
            stage_units_in_meters=1.0,
            physics_prim_path="/World/Environment/Physics/Scene",
        )
        world.get_physics_context().set_gravity(-9.81)

        # Potentially add a ground plane to the stage
        if add_ground_plane:
            world.scene.add(
                isaac_core.objects.ground_plane.GroundPlane(
                    prim_path="/World/Environment/Ground", name="ground", z_position=0.0
                )
            )

        # Potentially add a default light of type "DomeLight"
        if add_default_light:
            isaac_core.utils.prims.create_prim(
                "/World/Environment/DefaultLight",
                "DomeLight",
                attributes={"inputs:intensity": 1000.0},
                orientation=isaac_core.utils.rotations.euler_angles_to_quat(
                    (0, 45, 90), degrees=True
                ),
            )

        return world

    def load_usd(
        self,
        prim_path: str,
        absolute_usd_path: Optional[str] = None,
        relative_usd_path: Optional[str] = None,
        prim_type: str = "Xform",
        scale: list[float] = [1.0, 1.0, 1.0],
        position: Optional[list[float]] = None,
        translation: Optional[list[float]] = None,
        orientation: Optional[list[float]] = None,
    ) -> str:
        """Load a prim from a usd file into the scene.

        :param prim_path: Desired prim path for the prim in the stage.
        :type prim_path: str
        :param absolute_usd_path: Absolute path to the USD file in the docker container. Defaults to None.
        :type absolute_usd_path: Optional[str], optional
        :param relative_usd_path: Path to the USD file relative to the data root. Defaults to None.
        :type relative_usd_path: Optional[str], optional
        :param prim_type: The type of prim to add. Defaults to "Xform".
        :type prim_type: str, optional
        :param scale: The scaling factor in x, y, and z direction for the prim. Defaults to [1.0, 1.0, 1.0].
        :type scale: list[float], optional
        :param position: The position (xyz) at which to create the prim. If position and translation parameters are
         both None, position is used and set to [0.0, 0.0, 0.0].
        :type position: Optional[list[float]], optional
        :param translation: The translation (xyz) of the prim relative to its parent prim. If translation and position
         parameters are both none, position is used and set to [0.0, 0.0, 0.0].
        :type translation: Optional[list[float]], optional
        :param orientation: The orientation for the prim. If len 3: intrinsic Euler angles (xy'z'') in deg are expected,
         if len 4: quaternion (wxyz) is expected. Defaults to quaternion [1.0, 0.0, 0.0, 0.0].
        :type orientation: Optional[list[float]], optional
        """
        if position is None and translation is None:
            position = [0.0, 0.0, 0.0]

        # Get the USD file path
        if absolute_usd_path is not None:
            usd_path = absolute_usd_path
        elif relative_usd_path is not None:
            usd_path = f"{self.data_root}/{relative_usd_path}"
        else:
            carb.log_warn(
                "Must define either absolute or relative USD path for load_usd function."
            )
            return

        # Get the orientation as quaternion
        orientation = transform_utils.get_quat_from_orientation(orientation)
        orientation = orientation if orientation is not None else [1.0, 0.0, 0.0, 0.0]

        # Get a unique prim path starting from the given prim path
        prim_path = isaac_core.utils.string.find_unique_string_name(
            initial_name=prim_path,
            is_unique_fn=lambda x: not isaac_core.utils.prims.is_prim_path_valid(x),
        )

        # Create the prim
        isaac_core.utils.prims.create_prim(
            prim_path,
            prim_type,
            usd_path=usd_path,
            scale=scale,
            position=position,
            translation=translation,
            orientation=orientation,
        )

        return prim_path

    def make_prim_rigid_body(
        self,
        prim_path: str,
        collision_approximation: str = "convexHull",
        kinematic=False,
    ) -> None:
        """Make the prim at the given path a rigid body and adds a collider to it.

        :param prim_path: Path to the prim in the stage
        :type prim_path: str
        :param collision_approximation: The collision approximation shape. Can be "none" meaning triangle mesh collision,
         "convexHull", "convexDecomposition", "meshSimplification", "convexMeshSimplification", "boundingCube",
         "boundingSphere", "sphereFill", and "sdf". If it is None, the "Collision Enabled" option will be set to False.
          Defaults to "convexHull".
        :type collision_approximation: str, optional
        :param kinematic: Whether to make the rigid body kinematic. Defaults to False.
        :type kinematic: bool, optional
        """
        # Make prim rigid body
        prim = self.stage.GetPrimAtPath(prim_path)
        physx_utils.setRigidBody(
            prim, approximationShape=collision_approximation, kinematic=kinematic
        )

        # Potentially disable collision
        if collision_approximation is None:
            self.enable_prim_collider(prim_path=prim_path, enable=False)

    def undo_prim_rigid_body(self, prim_path: str, keep_collider: bool = False) -> None:
        """Remove the rigid body properties (RigidBodyAPI) from the prim at the given path.

        :param prim_path: Path to the prim.
        :type prim_path: str
        :param keep_collider: Whether to keep the collider. Defaults to False.
        :type keep_collider: bool, optional
        """
        prim = self.stage.GetPrimAtPath(prim_path)
        if keep_collider:
            physx_utils.removePhysics(prim)
        else:
            physx_utils.removeRigidBody(prim)

    def add_prim_collider(
        self, prim_path: str, collision_approximation: str = "none"
    ) -> None:
        """Adds a collider to the prim at the given path.

        :param prim_path: Path to the prim in the stage.
        :type prim_path: str
        :param collision_approximation: The collision approximation shape. Can be "none" meaning triangle mesh collision,
         "convexHull", "convexDecomposition", "meshSimplification", "convexMeshSimplification", "boundingCube",
         "boundingSphere", "sphereFill", and "sdf". If it is None, the "Collision Enabled" option will be set to False.
          Defaults to "none".
        :type collision_approximation: str, optional
        """
        # Add collider to the prim
        prim = self.stage.GetPrimAtPath(prim_path)
        if prim.IsA(UsdGeom.Xformable):
            physx_utils.setColliderSubtree(
                prim, approximationShape=collision_approximation
            )
        else:
            physx_utils.setCollider(prim, approximationShape=collision_approximation)

        # Potentially disable collision
        if collision_approximation is None:
            self.enable_prim_collider(prim_path=prim_path, enable=False)

    def remove_prim_collider(self, prim_path: str) -> None:
        """Removes the collider (CollisionAPI) from the prim at the given path.

        :param prim_path: Path to the prim in the stage.
        :type prim_path: str
        """
        prim = self.stage.GetPrimAtPath(prim_path)
        prim_children = iter(Usd.PrimRange(prim))
        for p in prim_children:
            if p.HasAPI(UsdPhysics.CollisionAPI):
                physx_utils.removeCollider(prim)

    def enable_prim_collider(self, prim_path: str, enable: bool = True) -> None:
        """Enable or disable collision for the prim at the given prim path.

        :param prim_path: Path to the prim in the stage.
        :type prim_path: str
        :param enable: Whether to enable collision for the prim. If False, collision will be disabled.
         Defaults to True.
        :type enable: bool, optional
        """
        prim = self.stage.GetPrimAtPath(prim_path)
        prim_children = iter(Usd.PrimRange(prim))
        for p in prim_children:
            if p.HasAPI(UsdPhysics.CollisionAPI):
                p.GetAttribute("physics:collisionEnabled").Set(enable)

    def prim_apply_visual_material(
        self, prim, material, weaker_than_descendants: bool = True
    ) -> None:
        # prim and material: Union[str, Usd.Prim]
        if isinstance(prim, str):
            prim = self.stage.GetPrimAtPath(prim)
        if isinstance(material, str):
            material = self.stage.GetPrimAtPath(material)

        material_shade = UsdShade.Material(material)
        if weaker_than_descendants:
            descendants_token = UsdShade.Tokens.weakerThanDescendants
        else:
            descendants_token = UsdShade.Tokens.strongerThanDescendants
        UsdShade.MaterialBindingAPI(prim).Bind(material_shade, descendants_token)

    def run_until_closed(self):
        """Runs the simulation until it is closed."""
        while self.simulation.is_running():
            self.world.step(render=True)
            if self.world.is_playing():
                if self.world.current_time_step_index == 1:
                    self.world.reset()
        self.simulation.close()