Add benchmark metrics to simulator

omnigibson/envs/env_base.py

@@ -1,3 +1,4 @@
+import time
 from copy import deepcopy
 
 import gymnasium as gym
@@ -50,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]
 
@@ -122,7 +120,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()
 
@@ -533,6 +531,11 @@ 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."""
+
+        # record the start time
+        # record the start time of the simulation step in the beginning of the step
+        self._cur_sim_start_ts = time.perf_counter()
+
         # 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()
@@ -560,6 +563,7 @@ def _post_step(self, action):
 
         # Grab reward, done, and info, and populate with internal info
         reward, done, info = self.task.step(self, action)
+
         self._populate_info(info)
         info["obs_info"] = obs_info
 
@@ -581,6 +585,11 @@ def _post_step(self, action):
 
         # Increment step
         self._current_step += 1
+
+        # record end time
+        # record the end time of the simulation step in the end of the step
+        self._prev_sim_end_ts = time.perf_counter()
+
         return obs, reward, terminated, truncated, info
 
     def step(self, action):
@@ -631,9 +640,14 @@ def _reset_variables(self):
         """
         Reset bookkeeping variables for the next new episode.
         """
+
         self._current_episode += 1
         self._current_step = 0
 
+        # reset the start and end time of the simulation step
+        self._prev_sim_end_ts = None
+        self._cur_sim_start_ts = None
+
     def reset(self, get_obs=True, **kwargs):
         """
         Reset episode.
@@ -690,6 +704,22 @@ def reset(self, get_obs=True, **kwargs):
 
             return obs, {}
 
+    @property
+    def last_step_wall_time(self):
+        """
+        Returns:
+            int: return the amount of wall time the last simulation step took
+        """
+
+        # return 0 if the simulation has not started yet
+        if not self._prev_sim_end_ts or not self._cur_sim_start_ts:
+            return 0
+
+        assert (
+            self._prev_sim_end_ts < self._cur_sim_start_ts
+        ), "end time from the previous iteration must be less than the start time of the current iteration"
+        return self._cur_sim_start_ts - self._prev_sim_end_ts
+
     @property
     def episode_steps(self):
         """

omnigibson/metrics/__init__.py

@@ -0,0 +1,5 @@
+from omnigibson.metrics.energy_metric import EnergyMetric
+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

omnigibson/metrics/energy_metric.py

@@ -0,0 +1,64 @@
+import numpy as np
+
+from omnigibson.metrics.metrics_base import BaseMetric
+
+
+class EnergyMetric(BaseMetric):
+    """
+    Energy Metric
+
+    Measures displacement * mass for every link
+
+    Args:
+        measure_work: If true, measure beginning and end delta rather than step by step delta
+    """
+
+    def __init__(self, measure_work=False):
+
+        # Run super
+        super().__init__()
+        self.state_cache = None
+        self.link_masses = {}
+
+        # parameter for checking if this is a work or energy metric. If true, measure work done, else measure energy
+        self.measure_work = measure_work
+
+    def _step(self, task, env, action):
+
+        # calculate the current pose of all object links
+        new_state_cache = {}
+        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)
+
+        # 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
+
+            for obj in env.scene.objects:
+                for link_name, link in obj._links.items():
+                    self.link_masses[link_name] = link.mass
+
+            return 0.0
+
+        # calculate the energy spent from the previous state to the current state
+        work_metric = 0.0
+        for linkname, posrot in new_state_cache.items():
+
+            # TODO: this computation is very slow, consider using a more efficient method
+            # TODO: this method needs to be updated to account for object addition and removal
+            posrot2 = self.state_cache[linkname]
+            work_metric += np.linalg.norm(posrot[0] - posrot2[0]) * self.link_masses[linkname]
+
+        # if measuring energy, update the state cache
+        if not self.measure_work:
+            self.state_cache = new_state_cache
+
+        # update the metric accordingly, either set the work done (measuring work) or add it to previous energy spent (measuring energy)
+        self._metric = work_metric if self.measure_work else (self._metric + work_metric)
+        return self._metric
+
+    def reset(self, task, env):
+        super().reset(task, env)
+        self.state_cache = None

omnigibson/metrics/metrics_base.py

@@ -0,0 +1,59 @@
+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
+        self._metric = 0
+
+    @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 = 0

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):
+        # Run super
+        super().__init__()
+
+    def _step(self, task, env, action):
+        self._metric += 1
+        return self._metric
+
+    def reset(self, task, env):
+        super().reset(task, env)

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):
+        # Run super
+        super().__init__()
+
+    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 self._metric
+
+    def reset(self, task, env):
+        super().reset(task, env)

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):
+        # Run super
+        super().__init__()
+
+    def _step(self, task, env, action):
+        self._metric += env.last_step_wall_time
+        return self._metric
+
+    def reset(self, task, env):
+        super().reset(task, env)

omnigibson/objects/controllable_object.py

@@ -576,7 +576,7 @@ def get_control_dict(self):
         # TODO: Move gravity force computation dummy to this class instead of BaseRobot
         fcns["gravity_force"] = lambda: (
             ControllableObjectViewAPI.get_generalized_gravity_forces(self.articulation_root_path)
-            if self.fixed_base
+            if self.fixed_base or self._dummy is None
             else ControllableObjectViewAPI.get_generalized_gravity_forces(self._dummy.articulation_root_path)
         )
         fcns["cc_force"] = lambda: ControllableObjectViewAPI.get_coriolis_and_centrifugal_forces(

omnigibson/tasks/behavior_task.py

@@ -10,6 +10,7 @@
     get_natural_goal_conditions,
     get_natural_initial_conditions,
     get_object_scope,
+    get_reward,
 )
 
 import omnigibson as og
@@ -285,12 +286,8 @@ def get_potential(self, env):
         Returns:
             float: Computed potential
         """
-        # Evaluate the first ground goal state option as the potential
-        _, satisfied_predicates = evaluate_goal_conditions(self.ground_goal_state_options[0])
-        success_score = len(satisfied_predicates["satisfied"]) / (
-            len(satisfied_predicates["satisfied"]) + len(satisfied_predicates["unsatisfied"])
-        )
-        return -success_score
+        self.success_score = get_reward(self.ground_goal_state_options)
+        return -self.success_score
 
     def initialize_activity(self, env):
         """
@@ -380,7 +377,9 @@ def _get_obs(self, env):
         low_dim_obs = dict()
 
         # Batch rpy calculations for much better efficiency
-        objs_exist = {obj: obj.exists for obj in self.object_scope.values() if not obj.is_system}
+        objs_exist = {
+            obj: obj.exists for obj in self.object_scope.values() if not obj.is_system and obj.states[Pose]._initialized
+        }
         objs_rpy = T.quat2euler(
             np.array(
                 [