diff --git a/omnigibson/configs/fetch_behavior.yaml b/omnigibson/configs/fetch_behavior.yaml index cced80efb..a875878c2 100644 --- a/omnigibson/configs/fetch_behavior.yaml +++ b/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 + diff --git a/omnigibson/envs/env_base.py b/omnigibson/envs/env_base.py index d84330be6..8fcc5734f 100644 --- a/omnigibson/envs/env_base.py +++ b/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): diff --git a/omnigibson/metrics/__init__.py b/omnigibson/metrics/__init__.py new file mode 100644 index 000000000..a8a5328c3 --- /dev/null +++ b/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 diff --git a/omnigibson/metrics/metrics_base.py b/omnigibson/metrics/metrics_base.py new file mode 100644 index 000000000..387538d96 --- /dev/null +++ b/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 diff --git a/omnigibson/metrics/step_metric.py b/omnigibson/metrics/step_metric.py new file mode 100644 index 000000000..88bfa6768 --- /dev/null +++ b/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 diff --git a/omnigibson/metrics/task_success_metric.py b/omnigibson/metrics/task_success_metric.py new file mode 100644 index 000000000..80e37248f --- /dev/null +++ b/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 diff --git a/omnigibson/metrics/wall_time_metric.py b/omnigibson/metrics/wall_time_metric.py new file mode 100644 index 000000000..a0d926563 --- /dev/null +++ b/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 diff --git a/omnigibson/metrics/work_energy_metric.py b/omnigibson/metrics/work_energy_metric.py new file mode 100644 index 000000000..69a78c57c --- /dev/null +++ b/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 = {} diff --git a/omnigibson/objects/controllable_object.py b/omnigibson/objects/controllable_object.py index 0ea7f3eed..6f78795ae 100644 --- a/omnigibson/objects/controllable_object.py +++ b/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__) diff --git a/omnigibson/prims/rigid_prim.py b/omnigibson/prims/rigid_prim.py index d58a45e1e..301bdbb10 100644 --- a/omnigibson/prims/rigid_prim.py +++ b/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): """ diff --git a/omnigibson/tasks/behavior_task.py b/omnigibson/tasks/behavior_task.py index 3c4da2195..d1fb00d61 100644 --- a/omnigibson/tasks/behavior_task.py +++ b/omnigibson/tasks/behavior_task.py @@ -10,11 +10,13 @@ get_natural_goal_conditions, get_natural_initial_conditions, get_object_scope, + get_reward, ) import omnigibson as og import omnigibson.utils.transform_utils as T from omnigibson.macros import gm +from omnigibson.metrics import StepMetric, TaskSuccessMetric, WallTimeMetric, WorkEnergyMetric from omnigibson.object_states import Pose from omnigibson.reward_functions.potential_reward import PotentialReward from omnigibson.robots.robot_base import BaseRobot @@ -66,6 +68,7 @@ def __init__( highlight_task_relevant_objects=False, termination_config=None, reward_config=None, + metric_config=None, ): # Make sure object states are enabled assert gm.ENABLE_OBJECT_STATES, "Must set gm.ENABLE_OBJECT_STATES=True in order to use BehaviorTask!" @@ -120,7 +123,9 @@ def __init__( ) # Run super init - super().__init__(termination_config=termination_config, reward_config=reward_config) + super().__init__( + termination_config=termination_config, reward_config=reward_config, metric_config=metric_config + ) @classmethod def get_cached_activity_scene_filename( @@ -189,6 +194,21 @@ def _create_reward_functions(self): return rewards + def _create_metric_functions(self): + """ + Creates the metric functions in the environment + + Returns: + dict of BaseRewardFunction: Metric functions created for Behavior task + """ + metrics = dict() + + metrics["steps"] = StepMetric() + metrics["task_success"] = TaskSuccessMetric() + metrics["wall_time"] = WallTimeMetric() + metrics["energy_work"] = WorkEnergyMetric(metric_config=self._metric_config) + return metrics + def _load(self, env): # Initialize the current activity success, self.feedback = self.initialize_activity(env=env) @@ -280,12 +300,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): """ @@ -565,3 +581,15 @@ def default_reward_config(cls): return { "r_potential": 1.0, } + + @classproperty + def default_metric_config(cls): + return { + "step": 0.05, + "task_success": 0.6, + "wall_time": 0.05, + "work": 0.2, + "energy": 0.2, + "rotation": 1, + "translation": 1, + } diff --git a/omnigibson/tasks/dummy_task.py b/omnigibson/tasks/dummy_task.py index 16a3fcd11..efbe9b53d 100644 --- a/omnigibson/tasks/dummy_task.py +++ b/omnigibson/tasks/dummy_task.py @@ -21,6 +21,10 @@ def _create_reward_functions(self): # Do nothing return dict() + def _create_metric_functions(self): + # Do nothing + return dict() + def _get_obs(self, env): # No task-specific obs of any kind return dict(), dict() @@ -43,3 +47,8 @@ def default_termination_config(cls): def default_reward_config(cls): # Empty dict return {} + + @classproperty + def default_metric_config(cls): + # Empty dict + return {} diff --git a/omnigibson/tasks/grasp_task.py b/omnigibson/tasks/grasp_task.py index 081690a73..782c755c2 100644 --- a/omnigibson/tasks/grasp_task.py +++ b/omnigibson/tasks/grasp_task.py @@ -78,6 +78,10 @@ def _create_reward_functions(self): rewards["grasp"] = GraspReward(self.obj_name, **self._reward_config) return rewards + def _create_metric_functions(self): + # No metric functions + return dict() + def _reset_agent(self, env): robot = env.robots[0] robot.release_grasp_immediately() @@ -237,3 +241,8 @@ def default_reward_config(cls): "eef_orientation_penalty_coef": 0.001, "regularization_coef": 0.01, } + + @classproperty + def default_metric_config(cls): + # Empty dict + return {} diff --git a/omnigibson/tasks/point_navigation_task.py b/omnigibson/tasks/point_navigation_task.py index 514bc56fb..1f96add64 100644 --- a/omnigibson/tasks/point_navigation_task.py +++ b/omnigibson/tasks/point_navigation_task.py @@ -148,6 +148,10 @@ def _create_reward_functions(self): return rewards + def _create_metric_functions(self): + # No metrics for this task + return dict() + def _load(self, env): # Load visualization self._load_visualization_markers(env=env) @@ -498,3 +502,8 @@ def default_reward_config(cls): "r_collision": 0.1, "r_pointgoal": 10.0, } + + @classproperty + def default_metric_config(cls): + # Empty dict + return {} diff --git a/omnigibson/tasks/point_reaching_task.py b/omnigibson/tasks/point_reaching_task.py index 743221faf..4948a2ee3 100644 --- a/omnigibson/tasks/point_reaching_task.py +++ b/omnigibson/tasks/point_reaching_task.py @@ -104,6 +104,10 @@ def _create_termination_conditions(self): return terminations + def _create_metric_functions(self): + # No metrics for this task + return dict() + def _sample_initial_pose_and_goal_pos(self, env, max_trials=100): # Run super first initial_pos, initial_ori, goal_pos = super()._sample_initial_pose_and_goal_pos(env=env, max_trials=max_trials) diff --git a/omnigibson/tasks/task_base.py b/omnigibson/tasks/task_base.py index 9801f2fa8..880b50983 100644 --- a/omnigibson/tasks/task_base.py +++ b/omnigibson/tasks/task_base.py @@ -1,3 +1,4 @@ +import time from abc import ABCMeta, abstractmethod from copy import deepcopy @@ -27,12 +28,13 @@ class BaseTask(GymObservable, Registerable, metaclass=ABCMeta): in with the default config. See cls.default_reward_config for default values used """ - def __init__(self, termination_config=None, reward_config=None): + def __init__(self, termination_config=None, reward_config=None, metric_config=None): # Make sure configs are dictionaries termination_config = dict() if termination_config is None else termination_config reward_config = dict() if reward_config is None else reward_config + metric_config = dict() if metric_config is None else metric_config - # Sanity check termination and reward conditions -- any keys found in the inputted config but NOT + # Sanity check termination, reward, and metric conditions -- any keys found in the inputted config but NOT # found in the default config should raise an error unknown_termination_keys = set(termination_config.keys()) - set(self.default_termination_config.keys()) assert ( @@ -40,20 +42,26 @@ def __init__(self, termination_config=None, reward_config=None): ), f"Got unknown termination config keys inputted: {unknown_termination_keys}" unknown_reward_keys = set(reward_config.keys()) - set(self.default_reward_config.keys()) assert len(unknown_reward_keys) == 0, f"Got unknown reward config keys inputted: {unknown_reward_keys}" + unknown_metric_keys = set(metric_config.keys()) - set(self.default_metric_config.keys()) + assert len(unknown_metric_keys) == 0, f"Got unknown metric config keys inputted: {unknown_metric_keys}" # Combine with defaults and store internally self._termination_config = self.default_termination_config self._termination_config.update(termination_config) self._reward_config = self.default_reward_config self._reward_config.update(reward_config) + self._metric_config = self.default_metric_config + self._metric_config.update(metric_config) # Generate reward and termination functions self._termination_conditions = self._create_termination_conditions() self._reward_functions = self._create_reward_functions() + self._metric_functions = self._create_metric_functions() # Store other internal vars that will be populated at runtime self._loaded = False self._reward = None + self._metrics = 0 self._done = None self._success = None self._info = None @@ -179,6 +187,17 @@ def _create_reward_functions(self): """ raise NotImplementedError() + def _create_metric_functions(self): + """ + Creates the metric functions in the environment + + Returns: + dict of BaseRewardFunction: Metric functions created for this task + """ + + # The metric functions are individually configured in the task class + raise NotImplementedError() + def _reset_scene(self, env): """ Task-specific scene reset. Default is the normal scene reset @@ -207,10 +226,15 @@ def _reset_variables(self, env): """ # By default, reset reward, done, and info self._reward = None + self._metrics = 0 self._done = False self._success = False self._info = None + # 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, env): """ Resets this task in the environment @@ -228,6 +252,8 @@ def reset(self, env): termination_condition.reset(self, env) for reward_function in self._reward_functions.values(): reward_function.reset(self, env) + for metric_function in self._metric_functions.values(): + metric_function.reset(self, env) def _step_termination(self, env, action, info=None): """ @@ -296,6 +322,31 @@ def _step_reward(self, env, action, info=None): return total_reward, total_info + def _step_metrics(self, env, action): + """ + Step and aggregate metric functions + + Args: + env (Environment): Environment instance + action (n-array): 1D flattened array of actions executed by all agents in the environment + + Returns: + - the break down of the metric scores and the metric info + """ + + # We'll also store individual metric split + breakdown_dict = dict() + metric_score = 0 + + for metric_function in self._metric_functions.values(): + metric = metric_function.step(self, env, action) + breakdown_dict.update(metric) + + for metric_name, metric_score in breakdown_dict.items(): + metric_score += metric_score * self._metric_config[metric_name] + + return metric_score, breakdown_dict + @abstractmethod def _get_obs(self, env): """ @@ -355,22 +406,33 @@ def step(self, env, action): # Make sure we're initialized assert self._loaded, "Task must be loaded using load() before calling step()!" - # We calculate termination conditions first and then rewards # (since some rewards can rely on termination conditions to update) done, done_info = self._step_termination(env=env, action=action) reward, reward_info = self._step_reward(env=env, action=action) + metric_score, metrics_info = self._step_metrics(env=env, action=action) # Update the internal state of this task self._reward = reward + self._metrics = metric_score self._done = done self._success = done_info["success"] self._info = { + "metrics": metrics_info, "reward": reward_info, "done": done_info, } return self._reward, self._done, deepcopy(self._info) + def pre_step(self): + + # record the real start time of the simulation step + self._cur_sim_start_ts = time.perf_counter() + + def post_step(self): + # record the real end time of the simulation step + self._prev_sim_end_ts = time.perf_counter() + @property def name(self): """ @@ -415,6 +477,22 @@ def info(self): assert self._info is not None, "At least one step() must occur before info can be calculated!" return self._info + @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 + @classproperty def valid_scene_types(cls): """ @@ -446,6 +524,16 @@ def default_termination_config(cls): """ raise NotImplementedError() + @classproperty + def default_metric_config(cls): + """ + Returns: + dict: Default metric configuration for this class. Should include any kwargs necessary for + any of the metric classes. Note: this default config should be fully verbose -- any keys + inputted in the constructor but NOT found in this default config will raise an error! + """ + raise NotImplementedError() + @classproperty def _do_not_register_classes(cls): # Don't register this class since it's an abstract template diff --git a/omnigibson/termination_conditions/predicate_goal.py b/omnigibson/termination_conditions/predicate_goal.py index 1e11cf0cd..75f690ce0 100644 --- a/omnigibson/termination_conditions/predicate_goal.py +++ b/omnigibson/termination_conditions/predicate_goal.py @@ -44,3 +44,11 @@ def goal_status(self): of the predicates matching either of those conditions """ return self._goal_status + + @property + def partial_success(self): + """ + Returns: + bool: returns true if partial success is supported, false otherwise + """ + return True diff --git a/omnigibson/termination_conditions/termination_condition_base.py b/omnigibson/termination_conditions/termination_condition_base.py index bd91049ce..e62c250ad 100644 --- a/omnigibson/termination_conditions/termination_condition_base.py +++ b/omnigibson/termination_conditions/termination_condition_base.py @@ -93,6 +93,14 @@ def success(self): assert self._done is not None, "At least one step() must occur before success can be calculated!" return self._done and self._terminate_is_success + @property + def partial_success(self): + """ + Returns: + bool: returns true if partial success is supported, false otherwise + """ + return False + @classproperty def _terminate_is_success(cls): """ diff --git a/tests/test_metric_functions.py b/tests/test_metric_functions.py new file mode 100644 index 000000000..df57da199 --- /dev/null +++ b/tests/test_metric_functions.py @@ -0,0 +1,156 @@ +import math +import os +from typing import OrderedDict + +import torch as th +import yaml +from bddl.condition_evaluation import evaluate_state + +import omnigibson as og +import omnigibson.utils.transform_utils as T +from omnigibson.macros import gm +from omnigibson.objects import DatasetObject +from omnigibson.tasks.behavior_task import BehaviorTask + +# Make sure object states are enabled +gm.ENABLE_OBJECT_STATES = True +gm.USE_GPU_DYNAMICS = True + + +def setup_env(): + + # Generates a BEHAVIOR Task environment using presampled cache + config_filename = os.path.join(og.example_config_path, "fetch_behavior.yaml") + cfg = yaml.load(open(config_filename, "r"), Loader=yaml.FullLoader) + cfg["task"]["online_object_sampling"] = False + + if og.sim is not None: + # Make sure sim is stopped + og.sim.stop() + + # Load the environment + env = og.Environment(configs=cfg) + env.reset() + return env + + +env = setup_env() + + +def test_behavior_task_work_metric(): + + # perform a step with no action + action = env.action_space.sample() + action["robot0"] = th.zeros_like(th.tensor(action["robot0"])) + + state, reward, terminated, truncated, info = env.step(action) + + metrics = info["metrics"] + + assert isinstance(metrics, dict) + + # assert that one step is taken + assert metrics["step"] == 1 + + # cache the initial position and orientation of the robot + position, orientation = env.robots[0].get_position_orientation() + + # move the robot to a new position and orientation, and then back to the initial position and orientation + env.robots[0].set_position_orientation([10, 10, 0]) + env.step(action) + + env.robots[0].set_position_orientation(position, orientation) + state, reward, terminated, truncated, info = env.step(action) + metrics = info["metrics"] + + # since simulator is running, the work done is non-zero due to random link movements. Assert that the work / robot mass is below a threshold + + # calculate robot mass + robot_mass = 0 + for link in env.robots[0].links.values(): + robot_mass += link.mass + + assert math.isclose(metrics["work"] / robot_mass, 0, abs_tol=1e-1) + env.reset() + + +def test_behavior_task_energy_metric(): + + # perform a step with no action + action = env.action_space.sample() + action["robot0"] = th.zeros_like(th.tensor(action["robot0"])) + env.step(action) + + # cache the initial position and orientation of the robot + position, orientation = env.robots[0].get_position_orientation() + + # apply shift to the robot + shift_position, shift_orientation = th.tensor([10, 10, 0], dtype=th.float32), th.tensor( + [0.05, 0, 0, 0.1], dtype=th.float32 + ) + env.robots[0].set_position_orientation(shift_position, shift_orientation) + state, reward, terminated, truncated, info = env.step(action) + energy = info["metrics"]["energy"] + + # shift the robot back to the initial position and orientation + env.robots[0].set_position_orientation(position, orientation) + state, reward, terminated, truncated, info = env.step(action) + new_energy = info["metrics"]["energy"] + + # since simulator is running, the work done is non-zero due to random link movements. Assert that the work / robot mass is below a threshold + + # calculate robot mass + robot_mass = 0 + for link in env.robots[0].links.values(): + robot_mass += link.mass + + # assert that the total energy spent is twice as much as the energy spent applying the shift + assert math.isclose((2 * energy - new_energy) / robot_mass, 0, abs_tol=1e-2) + env.reset() + + +def test_behavior_task_object_addition_removal(): + + # perform a step with no action + action = env.action_space.sample() + action["robot0"] = th.zeros_like(th.tensor(action["robot0"])) + env.step(action) + + env.robots[0].set_position_orientation([10, 10, 0]) + state, reward, terminated, truncated, info = env.step(action) + metrics = info["metrics"] + + work, energy = metrics["work"], metrics["energy"] + + # add another apple to the environment + apple = DatasetObject( + name="apple_unique", + category="apple", + model="agveuv", + ) + + env.scene.add_object(apple) + # perform a step with no action + state, reward, terminated, truncated, info = env.step(action) + metrics = info["metrics"] + + add_work, add_energy = metrics["work"], metrics["energy"] + + # calculate robot mass + robot_mass = 0 + for link in env.robots[0].links.values(): + robot_mass += link.mass + + assert math.isclose((work - add_work) / robot_mass, 0, abs_tol=1e-1) + assert math.isclose((energy - add_energy) / robot_mass, 0, abs_tol=1e-1) + + env.scene.remove_object(obj=apple) + state, reward, terminated, truncated, info = env.step(action) + metrics = info["metrics"] + + remove_work, remove_energy = metrics["work"], metrics["energy"] + + assert math.isclose((add_work - remove_work) / robot_mass, 0, abs_tol=1e-1) + assert math.isclose((add_energy - remove_energy) / robot_mass, 0, abs_tol=1e-1) + + env.reset()