Add benchmark metrics to simulator

omnigibson/configs/fetch_behavior.yaml

@@ -1,6 +1,7 @@
 env:
   action_frequency: 60                  # (int): environment executes action at the action_frequency rate
   physics_frequency: 60                 # (int): physics frequency (1 / physics_timestep for physx)
+  rendering_frequency: 60               # (int): rendering frequency
   device: null                          # (None or str): specifies the device to be used if running on the gpu with torch backend
   automatic_reset: false                # (bool): whether to automatic reset after an episode finishes
   flatten_action_space: false           # (bool): whether to flatten the action space as a sinle 1D-array
@@ -86,3 +87,12 @@ task:
     max_steps: 500
   reward_config:
     r_potential: 1.0
+  metric_config:
+    step: 0.05
+    task_success: 0.6
+    wall_time: 0.05
+    work: 0.2
+    energy: 0.2
+    rotation: 1
+    translation: 1
+

omnigibson/envs/env_base.py

@@ -51,9 +51,6 @@ def __init__(self, configs, in_vec_env=False):
         self.render_mode = "rgb_array"
         self.metadata = {"render.modes": ["rgb_array"]}
 
-        # Store if we are part of a vec env
-        self.in_vec_env = in_vec_env
-
         # Convert config file(s) into a single parsed dict
         configs = configs if isinstance(configs, list) or isinstance(configs, tuple) else [configs]
 
@@ -124,7 +121,7 @@ def __init__(self, configs, in_vec_env=False):
         self.load()
 
         # If we are not in a vec env, we can play ourselves. Otherwise we wait for the vec env to play.
-        if not self.in_vec_env:
+        if not in_vec_env:
             og.sim.play()
             self.post_play_load()
 
@@ -541,6 +538,10 @@ def _populate_info(self, info):
 
     def _pre_step(self, action):
         """Apply the pre-sim-step part of an environment step, i.e. apply the robot actions."""
+
+        # perform necessary pre-step actions for the task
+        self.task.pre_step()
+
         # 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()
@@ -589,6 +590,11 @@ def _post_step(self, action):
 
         # Increment step
         self._current_step += 1
+
+        # record end time
+        # perform necessary post-step actions for the task
+        self.task.post_step()
+
         return obs, reward, terminated, truncated, info
 
     def step(self, action, n_render_iterations=1):

omnigibson/metrics/__init__.py

@@ -0,0 +1,5 @@
+from omnigibson.metrics.metrics_base import BaseMetric
+from omnigibson.metrics.step_metric import StepMetric
+from omnigibson.metrics.task_success_metric import TaskSuccessMetric
+from omnigibson.metrics.wall_time_metric import WallTimeMetric
+from omnigibson.metrics.work_energy_metric import WorkEnergyMetric

omnigibson/metrics/metrics_base.py

@@ -0,0 +1,61 @@
+from abc import ABCMeta, abstractmethod
+from copy import deepcopy
+
+
+# from omnigibson.utils.python_utils import Registerable, classproperty
+class BaseMetric:
+    """
+    Base Metric class
+    Metric-specific reset and step methods are implemented in subclasses
+    """
+
+    def __init__(self):
+
+        # Store internal vars that will be filled in at runtime. The is a metric dictionary computed metric subclasses
+        self.metric = None
+
+    @abstractmethod
+    def _step(self, task, env, action):
+        """
+        Step the metric function and compute the metric at the current timestep. Overwritten by subclasses.
+
+        Args:
+            task (BaseTask): Task instance
+            env (Environment): Environment instance
+            action (n-array): 1D flattened array of actions executed by all agents in the environment
+
+        Returns:
+            2-tuple:
+                - bool: computed metric
+                - dict: any metric-related information for this specific metric
+        """
+        raise NotImplementedError()
+
+    def step(self, task, env, action):
+        """
+        Step the metricfunction and compute the metric at the current timestep. Overwritten by subclasses
+
+        Args:
+            task (BaseTask): Task instance
+            env (Environment): Environment instance
+            action (n-array): 1D flattened array of actions executed by all agents in the environment
+
+        Returns:
+            2-tuple:
+                - bool: computed metric
+                - dict: any metric-related information for this specific metric
+        """
+
+        # Step internally and store output
+        self.metric = self._step(task=task, env=env, action=action)
+
+        # Return metric
+        return self.metric
+
+    def reset(self, task, env):
+        """
+        General metrics reset
+        """
+
+        # Reset internal vars
+        self.metric = None

omnigibson/metrics/step_metric.py

@@ -0,0 +1,19 @@
+from omnigibson.metrics.metrics_base import BaseMetric
+
+
+class StepMetric(BaseMetric):
+    """
+    Step Metric
+    Metric for each simulator step
+    """
+
+    def __init__(self):
+        # initialize step
+        self._metric = 0
+
+    def _step(self, task, env, action):
+        self._metric += 1
+        return {"step": self._metric}
+
+    def reset(self, task, env):
+        self._metric = 0

omnigibson/metrics/task_success_metric.py

@@ -0,0 +1,39 @@
+from omnigibson.metrics.metrics_base import BaseMetric
+
+
+class TaskSuccessMetric(BaseMetric):
+    """
+    TaskSuccessMetric
+    Metric for partial or full task success
+    """
+
+    def __init__(self):
+        # initialize task success metric
+        self._metric = 0
+
+    def _step(self, task, env, action):
+        successes = []
+        partial_success = 0
+
+        # Evaluate termination conditions
+        for termination_condition in task._termination_conditions.values():
+
+            # Check if partial success is supported, and if so, store the score (e.g. Behavior Task)
+            if termination_condition.partial_success:
+                partial_success = task.success_score
+
+            done, success = termination_condition.step(task, env, action)
+            successes.append(success)
+
+        # Calculate metric
+        if any(successes):
+            self._metric = 1.0
+        elif partial_success > 0:
+            self._metric = partial_success
+        else:
+            self._metric = 0.0
+
+        return {"task_success": self._metric}
+
+    def reset(self, task, env):
+        self._metric = 0

omnigibson/metrics/wall_time_metric.py

@@ -0,0 +1,19 @@
+from omnigibson.metrics.metrics_base import BaseMetric
+
+
+class WallTimeMetric(BaseMetric):
+    """
+    WallTimeMetric
+    Metric for wall time accumulated in policy steps
+    """
+
+    def __init__(self):
+        # initialize wall time metric
+        self._metric = 0
+
+    def _step(self, task, env, action):
+        self._metric += task.last_step_wall_time
+        return {"wall_time": self._metric}
+
+    def reset(self, task, env):
+        self._metric = 0

omnigibson/metrics/work_energy_metric.py

@@ -0,0 +1,112 @@
+import torch as th
+
+import omnigibson.utils.transform_utils as T
+from omnigibson.metrics.metrics_base import BaseMetric
+
+
+class WorkEnergyMetric(BaseMetric):
+    """
+    Work and Energy Metric
+    Measures displacement * mass for every link
+    """
+
+    def __init__(self, metric_config=None):
+
+        # Run work and energy metric initialization
+        self._work_metric = 0
+        self._energy_metric = 0
+        self.state_cache = None
+        self.prev_state_cache = None
+
+        # stores object mass and rotational inertia
+        self.link_info = {}
+        self.metric_config = metric_config
+
+    def _step(self, task, env, action):
+
+        # calculate the current pose of all object links
+        new_state_cache = {}
+        self.link_info = {}
+        for obj in env.scene.objects:
+            for link_name, link in obj._links.items():
+                pos, rot = link.get_position_orientation()
+                new_state_cache[link_name] = (pos, rot)
+
+                # compute this every step to account for object addition and removal
+                # store link mass and inertia for energy calculation
+                mass = link.mass
+                inertia = link.inertia
+                self.link_info[link_name] = (mass, inertia)
+
+        # if the state cache is empty, set it to the current state and return 0
+        if not self.state_cache:
+            self.state_cache = new_state_cache
+            self.prev_state_cache = new_state_cache
+            return {"work": 0, "energy": 0}
+
+        # calculate the energy spent from the previous state to the current state
+        work_metric = 0.0
+        energy_metric = 0.0
+
+        for linkname, new_posrot in new_state_cache.items():
+
+            # TODO: this computation is very slow, consider using a more efficient method
+
+            # check if the link is originally in the state cache, if not, skip it to account for object addition and removal
+            if linkname not in self.state_cache:
+                continue
+
+            init_posrot = self.state_cache[linkname]
+            mass, inertia = self.link_info[linkname]
+            work_metric = self.calculate_work_energy(init_posrot, new_posrot, mass, inertia)
+
+            # check if the link is in the prev_state_cache, if not, skip it to account for object addition and removal
+            if linkname not in self.prev_state_cache:
+                continue
+
+            if self.prev_state_cache is not None:
+                prev_posrot = self.prev_state_cache[linkname]
+                energy_metric = self.calculate_work_energy(prev_posrot, new_posrot, mass, inertia)
+
+        # update the prev_state cache for energy measurement
+        self.prev_state_cache = new_state_cache
+
+        # update the metric accordingly, set the work done (measuring work) and add energy spent this step to previous energy spent (measuring energy)
+        self._work_metric = work_metric
+        self._energy_metric += energy_metric
+        return {"work": self._work_metric, "energy": self._energy_metric}
+
+    def calculate_work_energy(self, new_posrot, curr_posrot, mass, inertia):
+        """
+        Calculate the work done and energy spent in a single step
+
+        Args:
+            new_posrot (tuple): new position and orientation of the link
+            curr_posrot (tuple): current position and orientation of the link.
+                                 use initial posrot for work metric, and previous posrot for energy metric
+            mass (float): mass of the link
+            inertia (th.Tensor): rotational inertia of the link
+
+        Returns:
+            float: work/energy spent in a single step between two states
+        """
+
+        position, orientation = T.relative_pose_transform(new_posrot[0], new_posrot[1], curr_posrot[0], curr_posrot[1])
+        orientation = T.quat2axisangle(orientation)
+
+        # formula for translational work metric: displacement * mass * translation weight coefficient
+        work_metric = th.norm(position) * mass * self.metric_config["translation"]
+
+        # formula for rotational work metric: rotation * moment of inertia * rotation weight coefficient
+        work_metric += sum(th.matmul(orientation, inertia)) * self.metric_config["rotation"]
+
+        return work_metric
+
+    def reset(self, task, env):
+
+        # reset both _work_metric and _energy_metric, and the state cache
+        self._work_metric = 0
+        self._energy_metric = 0
+        self.state_cache = None
+        self.prev_state_cache = None
+        self.link_info = {}

omnigibson/objects/controllable_object.py

@@ -15,7 +15,11 @@
 from omnigibson.utils.numpy_utils import NumpyTypes
 from omnigibson.utils.python_utils import CachedFunctions, assert_valid_key, merge_nested_dicts
 from omnigibson.utils.ui_utils import create_module_logger
-from omnigibson.utils.usd_utils import ControllableObjectViewAPI
+from omnigibson.utils.usd_utils import (
+    ControllableObjectViewAPI,
+    absolute_prim_path_to_scene_relative,
+    add_asset_to_stage,
+)
 
 # Create module logger
 log = create_module_logger(module_name=__name__)

omnigibson/prims/rigid_prim.py

@@ -490,6 +490,21 @@ def mass(self):
 
         return mass
 
+    @property
+    def inertia(self):
+        """
+        Returns:
+            np.ndarray: inertia tensor of the rigid body in kg m^2.
+        """
+
+        inertia = self._rigid_prim_view.get_inertias()
+
+        # if no inertia is set, return a zero array
+        if inertia is None:
+            return th.zeros((3, 3))
+
+        return inertia.reshape(3, 3)
+
     @mass.setter
     def mass(self, mass):
         """