Add profiling

.github/workflows/profiling.yml

@@ -0,0 +1,53 @@
+name: Profiling
+
+on: 
+  push:
+    branches:
+      - og-develop
+
+permissions:
+  # deployments permission to deploy GitHub pages website
+  deployments: write
+  # contents permission to update profiling contents in gh-pages branch
+  contents: write
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event_name == 'pull_request' && github.head_ref || github.sha }}
+  cancel-in-progress: true
+
+jobs:
+  profiling:
+    name: Speed Profiling
+    runs-on: [self-hosted, linux, gpu, dataset-enabled]
+
+    defaults:
+      run:
+        shell: micromamba run -n omnigibson /bin/bash -leo pipefail {0}
+
+    steps:
+      - name: Fix home
+        run: echo "HOME=/root" >> $GITHUB_ENV
+
+      - name: Checkout source
+        uses: actions/checkout@v3
+
+      - name: Install dev requirements
+        run: pip install -r requirements-dev.txt
+
+      - name: Install
+        run: pip install -e .
+
+      - name: Run performance benchmark
+        run: bash scripts/profiling.sh
+
+      - name: Store benchmark result
+        uses: benchmark-action/github-action-benchmark@v1
+        with:
+          tool: 'customSmallerIsBetter'
+          output-file-path: output.json
+          benchmark-data-dir-path: profiling
+          fail-on-alert: false
+          alert-threshold: '200%'
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+          comment-on-alert: true
+          auto-push: true

omnigibson/macros.py

@@ -55,6 +55,9 @@
 # CANNOT be set at runtime
 gm.GUI_VIEWPORT_ONLY = False
 
+# Whether to use the viewer camera or not
+gm.RENDER_VIEWER_CAMERA = True
+
 # Do not suppress known omni warnings / errors, and also put omnigibson in a debug state
 # This includes extra information for things such as object sampling, and also any debug
 # logging messages

omnigibson/objects/object_base.py

@@ -177,11 +177,15 @@ def _post_load(self):
         if "visible" in self._load_config and self._load_config["visible"] is not None:
             self.visible = self._load_config["visible"]
 
-        # First, remove any articulation root API that already exists at the object-level prim
+        # First, remove any articulation root API that already exists at the object-level or root link level prim
         if self._prim.HasAPI(lazy.pxr.UsdPhysics.ArticulationRootAPI):
             self._prim.RemoveAPI(lazy.pxr.UsdPhysics.ArticulationRootAPI)
             self._prim.RemoveAPI(lazy.pxr.PhysxSchema.PhysxArticulationAPI)
 
+        if self.root_prim.HasAPI(lazy.pxr.UsdPhysics.ArticulationRootAPI):
+            self.root_prim.RemoveAPI(lazy.pxr.UsdPhysics.ArticulationRootAPI)
+            self.root_prim.RemoveAPI(lazy.pxr.PhysxSchema.PhysxArticulationAPI)
+
         # Potentially add articulation root APIs and also set self collisions
         root_prim = None if self.articulation_root_path is None else lazy.omni.isaac.core.utils.prims.get_prim_at_path(self.articulation_root_path)
         if root_prim is not None:

omnigibson/prims/entity_prim.py

@@ -898,43 +898,51 @@ def get_relative_angular_velocity(self):
 
     def set_position_orientation(self, position=None, orientation=None):
         # Delegate to RigidPrim if we are not articulated
-        if self._articulation_view is None:
-            return self.root_link.set_position_orientation(position=position, orientation=orientation)
-
-        if position is not None:
-            position = np.asarray(position)[None, :]
-        if orientation is not None:
-            orientation = np.asarray(orientation)[None, [3, 0, 1, 2]]
-        self._articulation_view.set_world_poses(position, orientation)
-        BoundingBoxAPI.clear()
+        if og.sim.is_stopped():
+            XFormPrim.set_position_orientation(self, position=position, orientation=orientation)
+        elif self._articulation_view is None:
+            self.root_link.set_position_orientation(position=position, orientation=orientation)
+        else:
+            if position is not None:
+                position = np.asarray(position)[None, :]
+            if orientation is not None:
+                orientation = np.asarray(orientation)[None, [3, 0, 1, 2]]
+            self._articulation_view.set_world_poses(position, orientation)
+            BoundingBoxAPI.clear()
 
     def get_position_orientation(self):
+        if og.sim.is_stopped():
+            return XFormPrim.get_position_orientation(self)
         # Delegate to RigidPrim if we are not articulated
-        if self._articulation_view is None:
+        elif self._articulation_view is None:
             return self.root_link.get_position_orientation()
-
-        positions, orientations = self._articulation_view.get_world_poses()
-        return positions[0], orientations[0][[1, 2, 3, 0]]
+        else:
+            positions, orientations = self._articulation_view.get_world_poses()
+            return positions[0], orientations[0][[1, 2, 3, 0]]
 
     def set_local_pose(self, position=None, orientation=None):
+        if og.sim.is_stopped():
+            XFormPrim.set_local_pose(self, position=position, orientation=orientation)
         # Delegate to RigidPrim if we are not articulated
-        if self._articulation_view is None:
-            return self.root_link.set_local_pose(position=position, orientation=orientation)
-
-        if position is not None:
-            position = np.asarray(position)[None, :]
-        if orientation is not None:
-            orientation = np.asarray(orientation)[None, [3, 0, 1, 2]]
-        self._articulation_view.set_local_poses(position, orientation)
-        BoundingBoxAPI.clear()
+        elif self._articulation_view is None:
+            self.root_link.set_local_pose(position=position, orientation=orientation)
+        else:
+            if position is not None:
+                position = np.asarray(position)[None, :]
+            if orientation is not None:
+                orientation = np.asarray(orientation)[None, [3, 0, 1, 2]]
+            self._articulation_view.set_local_poses(position, orientation)
+            BoundingBoxAPI.clear()
 
     def get_local_pose(self):
+        if og.sim.is_stopped():
+            return XFormPrim.get_local_pose(self)
         # Delegate to RigidPrim if we are not articulated
-        if self._articulation_view is None:
+        elif self._articulation_view is None:
             return self.root_link.get_local_pose()
-
-        positions, orientations = self._articulation_view.get_local_poses()
-        return positions[0], orientations[0][[1, 2, 3, 0]]
+        else:
+            positions, orientations = self._articulation_view.get_local_poses()
+            return positions[0], orientations[0][[1, 2, 3, 0]]
 
     # TODO: Is the omni joint damping (used for driving motors) same as dissipative joint damping (what we had in pb)?
     @property

omnigibson/simulator.py

@@ -1269,6 +1269,7 @@ def _init_stage(
                 position=np.array(m.DEFAULT_VIEWER_CAMERA_POS),
                 orientation=np.array(m.DEFAULT_VIEWER_CAMERA_QUAT),
             )
+            self.viewer_visibility = gm.RENDER_VIEWER_CAMERA
 
         def close(self):
             """

omnigibson/systems/macro_particle_system.py

@@ -1127,6 +1127,9 @@ class MacroPhysicalParticleSystem(PhysicalParticleSystem, MacroParticleSystem):
     _particle_radius = None
     _particle_offset = None
 
+    has_particles = False
+    has_refreshed = False
+
     @classmethod
     def initialize(cls):
         # Run super method first
@@ -1214,6 +1217,12 @@ def get_particles_position_orientation(cls):
             pos, ori = np.array([]).reshape(0, 3), np.array([]).reshape(0, 4)
         return pos, ori
 
+    @classmethod
+    def update(cls):
+        if cls.has_particles and not cls.has_refreshed:
+            cls._refresh_particles_view()
+            cls.has_refreshed = True
+
     @classmethod
     def get_particles_local_pose(cls):
         return cls.get_particles_position_orientation()
@@ -1369,6 +1378,9 @@ def generate_particles(
         # Set the vels
         cls.set_particles_velocities(lin_vels=velocities, ang_vels=angular_velocities)
 
+        if not cls.has_particles:
+            cls.has_particles = True
+
     @classmethod
     def create(cls, name, create_particle_template, particle_density, scale=None, **kwargs):
         """

omnigibson/systems/system_base.py

@@ -1231,6 +1231,10 @@ def is_fluid_system(system_name):
 def get_system(system_name, force_active=True):
     # Make sure scene exists
     assert og.sim.scene is not None, "Cannot get systems until scene is imported!"
+    # Make sure prefixes preserve their double underscore
+    for prefix in SYSTEM_PREFIXES:
+        if f"{prefix}__" not in system_name:
+            system_name = system_name.replace(f"{prefix}_", f"{prefix}__")
     # If system_name is not in REGISTERED_SYSTEMS, create from metadata
     system = REGISTERED_SYSTEMS[system_name] if system_name in REGISTERED_SYSTEMS \
         else _create_system_from_metadata(system_name=system_name)

omnigibson/transition_rules.py

@@ -975,7 +975,7 @@ def _generate_conditions(cls):
     @classmethod
     def transition(cls, object_candidates):
         objs_to_remove = []
-
+        
         for diceable_obj in object_candidates["diceable"]:
             obj_category = diceable_obj.category
             # We expect all diced particle systems to follow the naming convention (cooked__)diced__<category>

omnigibson/utils/profiling_utils.py

@@ -0,0 +1,80 @@
+import gym
+import omnigibson as og
+import os
+import psutil
+from pynvml.smi import nvidia_smi
+from time import time
+
+
+class ProfilingEnv(og.Environment):
+    def step(self, action):
+        try:
+            start = time()
+            # If the action is not a dictionary, convert into a dictionary
+            if not isinstance(action, dict) and not isinstance(action, gym.spaces.Dict):
+                action_dict = dict()
+                idx = 0
+                for robot in self.robots:
+                    action_dim = robot.action_dim
+                    action_dict[robot.name] = action[idx: idx + action_dim]
+                    idx += action_dim
+            else:
+                # Our inputted action is the action dictionary
+                action_dict = action
+
+            # Iterate over all robots and apply actions
+            for robot in self.robots:
+                robot.apply_action(action_dict[robot.name])
+
+            # Run simulation step
+            sim_start = time()
+            if len(og.sim._objects_to_initialize) > 0:
+                og.sim.render()
+            super(type(og.sim), og.sim).step(render=True)
+            omni_time = (time() - sim_start) * 1e3
+
+            # Additionally run non physics things
+            og.sim._non_physics_step()
+
+            # Grab observations
+            obs = self.get_obs()
+
+            # Step the scene graph builder if necessary
+            if self._scene_graph_builder is not None:
+                self._scene_graph_builder.step(self.scene)
+
+            # Grab reward, done, and info, and populate with internal info
+            reward, done, info = self.task.step(self, action)
+            self._populate_info(info)
+
+            if done and self._automatic_reset:
+                # Add lost observation to our information dict, and reset
+                info["last_observation"] = obs
+                obs = self.reset()
+
+            # Increment step
+            self._current_step += 1
+
+            # collect profiling data
+            total_frame_time = (time() - start) * 1e3
+            og_time = total_frame_time - omni_time
+            # memory usage in GB
+            memory_usage = psutil.Process(os.getpid()).memory_info().rss / 1024 ** 3
+            # VRAM usage in GB
+            for gpu in nvidia_smi.getInstance().DeviceQuery()['gpu']:
+                found = False
+                for process in gpu['processes']:
+                    if process['pid'] == os.getpid():
+                        vram_usage = process['used_memory'] / 1024
+                        found = True
+                        break
+                if found:
+                    break
+
+            ret = [total_frame_time, omni_time, og_time, memory_usage, vram_usage]
+            if self._current_step % 100 == 0:
+                print("total time: {:.3f} ms, Omni time: {:.3f} ms, OG time: {:.3f} ms, memory: {:.3f} GB, vram: {:.3f} GB.".format(*ret))
+
+            return obs, reward, done, info, ret
+        except:
+            raise ValueError(f"Failed to execute environment step {self._current_step} in episode {self._current_episode}")

requirements-dev.txt

@@ -7,4 +7,5 @@ mkdocs-material
 mkdocs-material-extensions
 mkdocstrings[python]
 mkdocs-section-index
-mkdocs-literate-nav
+mkdocs-literate-nav
+pynvml