Add profiling

.github/workflows/profiling.yml

@@ -0,0 +1,53 @@
+name: Profiling
+
+on: 
+  pull_request:
+    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: true
+          alert-threshold: '200%'
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+          comment-on-alert: true
+          auto-push: true

omnigibson/systems/system_base.py

@@ -1233,6 +1233,10 @@ def is_physical_particle_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

@@ -799,9 +799,8 @@ def _generate_conditions(cls):
     @classmethod
     def transition(cls, object_candidates):
         objs_to_remove = []
-
         for diceable_obj in object_candidates["diceable"]:
-            system = get_system(f"diced_{diceable_obj.category}")
+            system = get_system(f"diced__{diceable_obj.category}")
             system.generate_particles_from_link(diceable_obj, diceable_obj.root_link, check_contact=False, use_visual_meshes=False)
 
             # Delete original object from stage.

omnigibson/utils/profiling_utils.py

@@ -0,0 +1,56 @@
+import gym
+import omnigibson as og
+
+from time import time
+from omnigibson.envs.env_base import Environment
+
+PROFILING_FIELDS = ["total time", "physics time", "render time", "non physics time", "get observation time", "task time", "action time"]
+
+class ProfilingEnv(Environment):
+    def step(self, action):
+        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])
+        action_end = time()
+        # Run simulation step
+        # Possibly force playing
+        for i in range(og.sim.n_physics_timesteps_per_render):
+            super(type(og.sim), og.sim).step(render=False)
+        physics_end = time()
+        og.sim.render()
+        render_end = time()
+        og.sim._non_physics_step()
+        non_physics_end = time()
+        # Grab observations
+        obs = self.get_obs()
+        obs_end = time()
+        # Grab reward, done, and info, and populate with internal info, then get scene state
+        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()
+        # save scene state
+        self.scene.dump_state(serialized=True)
+        # Increment step
+        self._current_step += 1
+        end = time()
+        ret = [end-start, physics_end-action_end, render_end-physics_end, non_physics_end-render_end, \
+            obs_end-non_physics_end, end-obs_end, action_end-start]
+        if self._current_step % 100 == 0:
+            print("total time: {:.3f} ms, physics time: {:.3f} ms, render time: {:.3f} ms, non physics time: {:.3f} ms, get obs time: {:.3f} ms, task time: {:.3f} ms,  action time: {:.3f} ms, ".format(*ret))
+        return obs, reward, done, info, ret

scripts/benchmark.html → scripts/profiling.html

@@ -3,7 +3,7 @@
     <head>
         <title>OmniGibson Profiling</title>
         <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/bootstrap@5.0.0/dist/css/bootstrap.min.css">
-        <link rel="stylesheet" href="benchmark.css">
+        <link rel="stylesheet" href="profiling.css">
         <script src="https://cdn.jsdelivr.net/npm/bootstrap@5.0.0/dist/js/bootstrap.bundle.min.js"></script>
         <script src="https://cdnjs.cloudflare.com/ajax/libs/jquery/3.3.1/jquery.min.js"></script>
         <script src="https://cdn.jsdelivr.net/npm/chart.js"></script>
@@ -129,6 +129,6 @@ <h2>Scenes</h2>
             <div class="spacer"></div>
         </footer>
         <script src="data.js"></script>
-        <script src="benchmark.js"></script>
+        <script src="profiling.js"></script>
     </body>
 </html>

scripts/profiling.sh

@@ -0,0 +1,16 @@
+# 1st batch: baselines
+python tests/benchmark/profiling.py                  # baseline (fastest config possible)
+python tests/benchmark/profiling.py -s Rs_int        # for vision research
+python tests/benchmark/profiling.py -s Rs_int -r 1   # for basic robotics research
+# python tests/benchmark/profiling.py -s Rs_int -r 3   # for multi-agent research
+
+# 2nd batch: compare different scenes
+# python tests/benchmark/profiling.py -r 1 -s house_single_floor
+# python tests/benchmark/profiling.py -r 1 -s grocery_store_cafe
+# python tests/benchmark/profiling.py -r 1 -s Pomaria_0_garden
+
+# 3rd batch: OG non-physics features
+python tests/benchmark/profiling.py -r 1 -s Rs_int -w             # fluids (water)
+python tests/benchmark/profiling.py -r 1 -s Rs_int -c             # soft body (cloth)
+# python tests/benchmark/profiling.py -r 1 -s Rs_int -p             # macro particle system (diced objects)
+# python tests/benchmark/profiling.py -r 1 -s Rs_int -w -c -p       # everything

tests/benchmark/profiling.py

@@ -0,0 +1,169 @@
+import os
+import argparse
+import json
+import omnigibson as og
+import numpy as np
+import omnigibson.utils.transform_utils as T
+
+from omnigibson.macros import gm
+from omnigibson.systems import get_system
+from omnigibson.object_states import Covered
+from omnigibson.utils.profiling_utils import ProfilingEnv
+from omnigibson.utils.constants import PrimType
+
+parser = argparse.ArgumentParser()
+
+parser.add_argument("-r", "--robot", type=int, default=0)
+parser.add_argument("-s", "--scene", default="")
+parser.add_argument("-f", "--flatcache", action='store_true')
+parser.add_argument("-c", "--cloth", action='store_true')
+parser.add_argument("-w", "--fluids", action='store_true')
+parser.add_argument("-p", "--macro_particle_system", action='store_true')
+
+PROFILING_FIELDS = ["Total frame time", "Physics step time", "Render step time", "Non-physics step time"]
+NUM_CLOTH = 5
+NUM_SLICE_OBJECT = 3
+
+SCENE_OFFSET = {
+    "": [0, 0],
+    "Rs_int": [0, 0],
+    "Pomaria_0_garden": [0.3, 0],
+    "grocery_store_cafe": [-3.5, 3.5],
+    "house_single_floor": [0, 0],
+}
+
+
+def main():
+    args = parser.parse_args()
+    # Modify flatcache, pathtracing, GPU, and object state settings
+    assert not (args.flatcache and args.cloth), "Flatcache cannot be used with softbody at the same time"
+    gm.ENABLE_FLATCACHE = args.flatcache
+    gm.ENABLE_HQ_RENDERING = args.fluids
+    gm.ENABLE_OBJECT_STATES = True
+    gm.ENABLE_TRANSITION_RULES = True
+    gm.USE_GPU_DYNAMICS = True
+
+    cfg = {
+        "env": {
+            "action_frequency": 60,
+            "physics_frequency": 300,
+        }
+    }
+    if args.robot > 0:
+        cfg["robots"] = []
+        for i in range(args.robot):
+            cfg["robots"].append({
+                "type": "Fetch",
+                "obs_modalities":"all",
+                "position": [-1.3 + 0.75 * i + SCENE_OFFSET[args.scene][0], 0.5 + SCENE_OFFSET[args.scene][1], 0],
+                "orientation": [0., 0., 0.7071, -0.7071]
+            })
+
+    if args.scene:
+        assert args.scene in SCENE_OFFSET, f"Scene {args.scene} not found in SCENE_OFFSET"
+        cfg["scene"] = {
+            "type": "InteractiveTraversableScene",
+            "scene_model": args.scene,
+        }
+    else:
+        cfg["scene"] = {"type": "Scene"}
+
+    cfg["objects"] = [{
+        "type": "DatasetObject",
+        "name": "table",
+        "category": "breakfast_table",
+        "model": "rjgmmy",
+        "scale": [0.75] * 3,
+        "position": [0.5 + SCENE_OFFSET[args.scene][0], -1 + SCENE_OFFSET[args.scene][1], 0.4],
+        "orientation": [0., 0., 0.7071, -0.7071]
+    }]
+
+    if args.cloth:
+        cfg["objects"].extend([{
+            "type": "DatasetObject",
+            "name": f"shirt_{n}",
+            "category": "t_shirt",
+            "model": "kvidcx",
+            "prim_type": PrimType.CLOTH,
+            "scale": [0.01] * 3,
+            "position": [-0.4, -1, 0.7 + n * 0.4],
+            "orientation": [0.7071, 0., 0.7071, 0.],
+        } for n in range(NUM_CLOTH)])
+
+    cfg["objects"].extend([{
+        "type": "DatasetObject",
+        "name": f"apple_{n}",
+        "category": "apple",
+        "model": "agveuv",
+        "scale": [1.5] * 3,
+        "position": [0.5 + SCENE_OFFSET[args.scene][0], -1.25 + SCENE_OFFSET[args.scene][1] + n * 0.2, 0.6],
+        "abilities": {"diceable": {}} if args.macro_particle_system else {}
+    } for n in range(NUM_SLICE_OBJECT)])
+    cfg["objects"].extend([{
+        "type": "DatasetObject",
+        "name": f"knife_{n}",
+        "category": "table_knife",
+        "model": "jxdfyy",
+        "scale": [2.5] * 3,
+        "position": [0.5 + SCENE_OFFSET[args.scene][0], -1.25 + SCENE_OFFSET[args.scene][1] + n * 0.2, 0.8],
+        "orientation": T.euler2quat([-np.pi / 2, 0, 0])
+    } for n in range(NUM_SLICE_OBJECT)])
+
+    env = ProfilingEnv(configs=cfg)
+    env.reset()
+
+    table = env.scene.object_registry("name", "table")
+    apples = [env.scene.object_registry("name", f"apple_{n}") for n in range(NUM_SLICE_OBJECT)]
+    knifes = [env.scene.object_registry("name", f"knife_{n}") for n in range(NUM_SLICE_OBJECT)]
+    for apple, knife in zip(apples, knifes):
+        knife.keep_still()
+        knife.set_position_orientation(
+            position=apple.get_position() + np.array([-0.15, 0.0, 0.2]),
+            orientation=T.euler2quat([-np.pi / 2, 0, 0]),
+        )
+    if args.fluids:
+        table.states[Covered].set_value(get_system("water"), True)
+
+    output, results = [], []
+
+    # Update the simulator's viewer camera's pose so it points towards the robot
+    og.sim.viewer_camera.set_position([SCENE_OFFSET[args.scene][0], -3 + SCENE_OFFSET[args.scene][1], 1])
+
+    for i in range(500):
+        if args.robot:
+            result = env.step(np.array([np.random.uniform(-0.3, 0.3, env.robots[i].action_dim) for i in range(args.robot)]).flatten())[4][:4]
+        else:
+            result = env.step(None)[4][:4]
+        results.append(result)
+
+    results = np.array(results)
+    for i, title in enumerate(PROFILING_FIELDS):
+        field = f"{args.scene}" if args.scene else "Empty scene"
+        if args.robot:
+            field += f", with {args.robot} Fetch"
+        if args.cloth:
+            field += ", cloth" 
+        if args.fluids:
+            field += ", fluids"
+        if args.macro_particle_system:
+            field += ", macro particles"
+        if args.flatcache:
+            field += ", flatcache on"
+        output.append({
+            "name": field,
+            "unit": "frame time (ms)",
+            "value": 1 / np.mean(results[-300:, i]),
+            "extra": [title, title]
+        })
+
+    ret = []
+    if os.path.exists("output.json"):
+        with open("output.json", "r") as f:
+            ret = json.load(f)
+    ret.extend(output)
+    with open("output.json", "w") as f:
+        json.dump(ret, f, indent=4)
+    og.shutdown()
+
+if __name__ == "__main__":
+    main()