diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index a4c1109b0..3de441444 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -19,6 +19,9 @@ jobs: strategy: fail-fast: false matrix: + device: + - cpu + - cuda test_file: - test_controllers - test_curobo @@ -29,8 +32,7 @@ jobs: - test_object_removal - test_object_states - test_primitives - - test_robot_states_flatcache - - test_robot_states_no_flatcache + - test_robot_states - test_robot_teleoperation - test_scene_graph - test_sensors @@ -59,21 +61,21 @@ jobs: - name: Run tests working-directory: omnigibson-src - run: pytest -s tests/${{ matrix.test_file }}.py --junitxml=${{ matrix.test_file }}.xml && cp ${{ matrix.test_file }}.xml ${GITHUB_WORKSPACE}/ + run: pytest -s tests/${{ matrix.test_file }}.py -k ${{ matrix.device }} --junitxml=${{ matrix.test_file }}-${{ matrix.device }}.xml && cp ${{ matrix.test_file }}-${{ matrix.device }}.xml ${GITHUB_WORKSPACE}/ continue-on-error: true - - name: Deploy artifact + - name: Deploy artifacts uses: actions/upload-artifact@v3 with: name: ${{ github.run_id }}-tests-${{ matrix.test_file }} - path: ${{ matrix.test_file }}.xml + path: ${{ matrix.test_file }}-${{ matrix.device }}.xml - name: Check for failures, errors, or missing XML run: | if [ ! -f ${{ matrix.test_file }}.xml ]; then echo "Error: XML file not found, probably due to segfault" exit 1 - elif grep -Eq 'failures="[1-9][0-9]*"|errors="[1-9][0-9]*"' ${{ matrix.test_file }}.xml; then + elif grep -Eq 'failures="[1-9][0-9]*"|errors="[1-9][0-9]*"' ${{ matrix.test_file }}-${{ matrix.device }}.xml; then echo "Error: Test failures or errors found" exit 1 else diff --git a/docs/miscellaneous/speed_optimization.md b/docs/miscellaneous/speed_optimization.md index 93e3e5604..fcd68f249 100644 --- a/docs/miscellaneous/speed_optimization.md +++ b/docs/miscellaneous/speed_optimization.md @@ -15,12 +15,7 @@ A lot of factors could affect the speed of OmniGibson. Here are a few of them: 1. `ENABLE_HQ_RENDERING`: While it is set to False by default, setting it to True will give us better rendering quality as well as other more advanced rendering features (e.g. isosurface for fluids), but with the cost of dragging down performance. -2. `USE_GPU_DYNAMICS`: setting it to True is required for more advanced features like particles and fluids, but it will lower the performance of OmniGibson. - -3. `RENDER_VIEWER_CAMERA`: viewer camera refers to the camera that shows the default viewport in OmniGibson, if you don't want to view the entire scene (e.g. during training), you can set this to False andit will boost OmniGibson performance. - -4. `ENABLE_FLATCACHE`: setting it to True will boost OmniGibson performance. - +2. `RENDER_VIEWER_CAMERA`: viewer camera refers to the camera that shows the default viewport in OmniGibson, if you don't want to view the entire scene (e.g. during training), you can set this to False andit will boost OmniGibson performance. ## Miscellaneous diff --git a/omnigibson/controllers/controller_base.py b/omnigibson/controllers/controller_base.py index 7bb6feeb8..6e65da572 100644 --- a/omnigibson/controllers/controller_base.py +++ b/omnigibson/controllers/controller_base.py @@ -4,6 +4,7 @@ import torch as th +import omnigibson as og from omnigibson.utils.python_utils import Recreatable, Registerable, Serializable, assert_valid_key, classproperty # Global dicts that will contain mappings @@ -100,19 +101,25 @@ def __init__( to the @control_limits entry corresponding to self.control_type """ # Store arguments + assert "has_limit" in control_limits, "Expected has_limit specified in control_limits, but does not exist." + assert type(dof_idx) == list, f"Expected dof_idx to be a list, got type {type(dof_idx)} instead." + self._dof_idx = dof_idx + # Store the indices in self.dof_idx that have control limits + self._limited_dof_indices = th.tensor( + [i for i, idx in enumerate(self.dof_idx) if control_limits["has_limit"][idx]], + dtype=th.long, + device=og.sim.device, + ) self._control_freq = control_freq - self._control_limits = {} + # Store control limits as a 3 x 2 x n tensor: [control_type][min/max][dof_idx] + self._control_limits = th.zeros((3, 2, len(dof_idx)), device=og.sim.device) for motor_type in {"position", "velocity", "effort"}: - if motor_type not in control_limits: - continue - - self._control_limits[ControlType.get_type(motor_type)] = [ - control_limits[motor_type][0], - control_limits[motor_type][1], - ] - assert "has_limit" in control_limits, "Expected has_limit specified in control_limits, but does not exist." - self._dof_has_limits = control_limits["has_limit"] - self._dof_idx = dof_idx.int() + self._control_limits[ControlType.get_type(motor_type)][0] = control_limits[motor_type][0][dof_idx].to( + device=og.sim.device + ) + self._control_limits[ControlType.get_type(motor_type)][1] = control_limits[motor_type][1][dof_idx].to( + device=og.sim.device + ) # Generate goal information self._goal_shapes = self._get_goal_shapes() @@ -133,8 +140,8 @@ def __init__( ) command_output_limits = ( ( - self._control_limits[self.control_type][0][self.dof_idx], - self._control_limits[self.control_type][1][self.dof_idx], + self._control_limits[self.control_type][0], + self._control_limits[self.control_type][1], ) if type(command_output_limits) == str and command_output_limits == "default" else command_output_limits @@ -169,7 +176,11 @@ def _preprocess_command(self, command): Array[float]: Processed command vector """ # Make sure command is a th.tensor - command = th.tensor([command], dtype=th.float32) if type(command) in {int, float} else command + command = ( + th.tensor([command], dtype=th.float32, device=og.sim.device) + if type(command) in {int, float} + else command.to(device=og.sim.device) + ) # We only clip and / or scale if self.command_input_limits exists if self._command_input_limits is not None: # Clip @@ -238,7 +249,9 @@ def update_goal(self, command, control_dict): ), f"Commands must be dimension {self.command_dim}, got dim {len(command)} instead." # Preprocess and run internal command - self._goal = self._update_goal(command=self._preprocess_command(command), control_dict=control_dict) + self._goal = self._update_goal( + command=self._preprocess_command(command.to(device=og.sim.device)), control_dict=control_dict + ) def _update_goal(self, command, control_dict): """ @@ -281,15 +294,13 @@ def clip_control(self, control): Array[float]: Clipped control signal """ clipped_control = control.clip( - self._control_limits[self.control_type][0][self.dof_idx], - self._control_limits[self.control_type][1][self.dof_idx], + self._control_limits[self.control_type][0], + self._control_limits[self.control_type][1], ) - idx = ( - self._dof_has_limits[self.dof_idx] - if self.control_type == ControlType.POSITION - else [True] * self.control_dim - ) - control[idx] = clipped_control[idx] + if self.control_type == ControlType.POSITION: + control[self._limited_dof_indices] = clipped_control[self._limited_dof_indices] + else: + control = clipped_control return control def step(self, control_dict): @@ -418,12 +429,12 @@ def nums2array(nums, dim): # Check if input is an Iterable, if so, we simply convert the input to th.tensor and return # Else, input is a single value, so we map to a numpy array of correct size and return return ( - nums + nums.to(device=og.sim.device) if isinstance(nums, th.Tensor) else ( - th.tensor(nums, dtype=th.float32) + th.tensor(nums, dtype=th.float32, device=og.sim.device) if isinstance(nums, Iterable) - else th.ones(dim, dtype=th.float32) * nums + else th.ones(dim, dtype=th.float32, device=og.sim.device) * nums ) ) @@ -511,7 +522,7 @@ def command_dim(self): def dof_idx(self): """ Returns: - Array[int]: DOF indices corresponding to the specific DOFs being controlled by this robot + th.Tensor[int]: DOF indices corresponding to the specific DOFs being controlled by this robot """ return self._dof_idx diff --git a/omnigibson/controllers/dd_controller.py b/omnigibson/controllers/dd_controller.py index 38207b417..41e76c996 100644 --- a/omnigibson/controllers/dd_controller.py +++ b/omnigibson/controllers/dd_controller.py @@ -1,5 +1,6 @@ import torch as th +import omnigibson as og from omnigibson.controllers import ControlType, LocomotionController @@ -50,8 +51,8 @@ def __init__( @control_limits velocity limits entry """ # Store internal variables - self._wheel_radius = wheel_radius - self._wheel_axle_halflength = wheel_axle_length / 2.0 + self._wheel_radius = th.tensor(wheel_radius, device=og.sim.device) + self._wheel_axle_halflength = th.tensor(wheel_axle_length / 2.0, device=og.sim.device) # If we're using default command output limits, map this to maximum linear / angular velocities if type(command_output_limits) == str and command_output_limits == "default": @@ -107,11 +108,11 @@ def compute_control(self, goal_dict, control_dict): right_wheel_joint_vel = (lin_vel + ang_vel * self._wheel_axle_halflength) / self._wheel_radius # Return desired velocities - return th.tensor([left_wheel_joint_vel, right_wheel_joint_vel]) + return th.stack([left_wheel_joint_vel, right_wheel_joint_vel]) def compute_no_op_goal(self, control_dict): # This is zero-vector, since we want zero linear / angular velocity - return dict(vel=th.zeros(2)) + return dict(vel=th.zeros(2, device=og.sim.device)) def _compute_no_op_action(self, control_dict): return th.zeros(2, dtype=th.float32) diff --git a/omnigibson/controllers/ik_controller.py b/omnigibson/controllers/ik_controller.py index fa0e18b2a..d9b8b2a2e 100644 --- a/omnigibson/controllers/ik_controller.py +++ b/omnigibson/controllers/ik_controller.py @@ -2,6 +2,7 @@ import torch as th +import omnigibson as og import omnigibson.utils.transform_utils as T from omnigibson.controllers import ControlType, ManipulationController from omnigibson.controllers.joint_controller import JointController @@ -331,8 +332,8 @@ def compute_control(self, goal_dict, control_dict): # Clip values to be within the joint limits target_joint_pos = target_joint_pos.clamp( - min=self._control_limits[ControlType.get_type("position")][0][self.dof_idx], - max=self._control_limits[ControlType.get_type("position")][1][self.dof_idx], + min=self._control_limits[ControlType.get_type("position")][0], + max=self._control_limits[ControlType.get_type("position")][1], ) # Optionally pass through smoothing filter for better stability diff --git a/omnigibson/controllers/joint_controller.py b/omnigibson/controllers/joint_controller.py index 4ee046cd6..9a5b960ff 100644 --- a/omnigibson/controllers/joint_controller.py +++ b/omnigibson/controllers/joint_controller.py @@ -2,6 +2,7 @@ import torch as th +import omnigibson as og import omnigibson.utils.transform_utils as T from omnigibson.controllers import ( ControlType, @@ -111,9 +112,13 @@ def __init__( assert pos_kp is None, "Cannot set pos_kp for JointController with motor_type=effort!" assert pos_damping_ratio is None, "Cannot set pos_damping_ratio for JointController with motor_type=effort!" assert vel_kp is None, "Cannot set vel_kp for JointController with motor_type=effort!" - self.pos_kp = pos_kp - self.pos_kd = None if pos_kp is None or pos_damping_ratio is None else 2 * math.sqrt(pos_kp) * pos_damping_ratio - self.vel_kp = vel_kp + self.pos_kp = th.tensor(pos_kp, device=og.sim.device) if pos_kp is not None else None + self.pos_kd = ( + None + if pos_kp is None or pos_damping_ratio is None + else th.tensor(2 * math.sqrt(pos_kp) * pos_damping_ratio, device=og.sim.device) + ) + self.vel_kp = th.tensor(vel_kp, device=og.sim.device) if vel_kp is not None else None self._use_impedances = use_impedances self._use_gravity_compensation = use_gravity_compensation self._use_cc_compensation = use_cc_compensation @@ -159,7 +164,10 @@ def _update_goal(self, command, control_dict): # Compute the final rotations in the quaternion space. _, end_quat = T.pose_transform( - th.zeros(3), T.euler2quat(delta_rots), th.zeros(3), T.euler2quat(start_rots) + th.zeros(3, dtype=th.float32, device=og.sim.device), + T.euler2quat(delta_rots), + th.zeros(3, dtype=th.float32, device=og.sim.device), + T.euler2quat(start_rots), ) end_rots = T.quat2euler(end_quat) @@ -172,8 +180,8 @@ def _update_goal(self, command, control_dict): # Clip the command based on the limits target = target.clip( - self._control_limits[ControlType.get_type(self._motor_type)][0][self.dof_idx], - self._control_limits[ControlType.get_type(self._motor_type)][1][self.dof_idx], + self._control_limits[ControlType.get_type(self._motor_type)][0], + self._control_limits[ControlType.get_type(self._motor_type)][1], ) return dict(target=target) @@ -238,7 +246,7 @@ def compute_no_op_goal(self, control_dict): target = control_dict[f"joint_{self._motor_type}"][self.dof_idx] else: # For velocity / effort, directly set to 0 - target = th.zeros(self.control_dim) + target = th.zeros(self.control_dim, device=og.sim.device) return dict(target=target) diff --git a/omnigibson/controllers/multi_finger_gripper_controller.py b/omnigibson/controllers/multi_finger_gripper_controller.py index 9eb70241f..564b2920a 100644 --- a/omnigibson/controllers/multi_finger_gripper_controller.py +++ b/omnigibson/controllers/multi_finger_gripper_controller.py @@ -166,13 +166,13 @@ def compute_control(self, goal_dict, control_dict): should_open = target[0] >= 0.0 if not self._inverted else target[0] > 0.0 if should_open: u = ( - self._control_limits[ControlType.get_type(self._motor_type)][1][self.dof_idx] + self._control_limits[ControlType.get_type(self._motor_type)][1] if self._open_qpos is None else self._open_qpos ) else: u = ( - self._control_limits[ControlType.get_type(self._motor_type)][0][self.dof_idx] + self._control_limits[ControlType.get_type(self._motor_type)][0] if self._closed_qpos is None else self._closed_qpos ) @@ -182,12 +182,8 @@ def compute_control(self, goal_dict, control_dict): # If we're near the joint limits and we're using velocity / torque control, we zero out the action if self._motor_type in {"velocity", "torque"}: - violate_upper_limit = ( - joint_pos > self._control_limits[ControlType.POSITION][1][self.dof_idx] - self._limit_tolerance - ) - violate_lower_limit = ( - joint_pos < self._control_limits[ControlType.POSITION][0][self.dof_idx] + self._limit_tolerance - ) + violate_upper_limit = joint_pos > self._control_limits[ControlType.POSITION][1] - self._limit_tolerance + violate_lower_limit = joint_pos < self._control_limits[ControlType.POSITION][0] + self._limit_tolerance violation = th.logical_or(violate_upper_limit * (u > 0), violate_lower_limit * (u < 0)) u *= ~violation @@ -241,8 +237,8 @@ def _update_grasping_state(self, control_dict): # Otherwise, the last control signal intends to "move" the gripper else: finger_vel = control_dict["joint_velocity"][self.dof_idx] - min_pos = self._control_limits[ControlType.POSITION][0][self.dof_idx] - max_pos = self._control_limits[ControlType.POSITION][1][self.dof_idx] + min_pos = self._control_limits[ControlType.POSITION][0] + max_pos = self._control_limits[ControlType.POSITION][1] # Make sure we don't have any invalid values (i.e.: fingers should be within the limits) finger_pos = th.clip(finger_pos, min_pos, max_pos) diff --git a/omnigibson/controllers/osc_controller.py b/omnigibson/controllers/osc_controller.py index 01515fc3e..c66f1f143 100644 --- a/omnigibson/controllers/osc_controller.py +++ b/omnigibson/controllers/osc_controller.py @@ -2,6 +2,7 @@ import torch as th +import omnigibson as og import omnigibson.utils.transform_utils as T from omnigibson.controllers import ControlType, ManipulationController from omnigibson.utils.control_utils import orientation_error @@ -144,9 +145,13 @@ def limiter(target_pos: Array[float], target_quat: Array[float], control_dict: D ) # Store gains - self.kp = nums2array(nums=kp, dim=6, dtype=th.float32) if kp is not None else None + self.kp = nums2array(nums=kp, dim=6, dtype=th.float32).to(og.sim.device) if kp is not None else None self.damping_ratio = damping_ratio - self.kp_null = nums2array(nums=kp_null, dim=control_dim, dtype=th.float32) if kp_null is not None else None + self.kp_null = ( + nums2array(nums=kp_null, dim=control_dim, dtype=th.float32).to(og.sim.device) + if kp_null is not None + else None + ) self.kd_null = 2 * th.sqrt(self.kp_null) if kp_null is not None else None # critically damped self.kp_limits = th.tensor(kp_limits, dtype=th.float32) self.damping_ratio_limits = th.tensor(damping_ratio_limits, dtype=th.float32) @@ -235,7 +240,8 @@ def limiter(target_pos: Array[float], target_quat: Array[float], control_dict: D self.decouple_pos_ori = decouple_pos_ori self.workspace_pose_limiter = workspace_pose_limiter self.task_name = task_name - self.reset_joint_pos = reset_joint_pos[dof_idx] + self.reset_joint_pos = reset_joint_pos[dof_idx].to(og.sim.device) + self._null_space_identity = th.eye(control_dim, dtype=th.float32, device=og.sim.device) # Other variables that will be filled in at runtime self._fixed_quat_target = None @@ -425,6 +431,7 @@ def compute_control(self, goal_dict, control_dict): decouple_pos_ori=self.decouple_pos_ori, base_lin_vel=base_lin_vel, base_ang_vel=base_ang_vel, + null_space_identity=self._null_space_identity, ).flatten() # Add gravity compensation @@ -506,6 +513,7 @@ def _compute_osc_torques( decouple_pos_ori: bool, base_lin_vel: th.Tensor, base_ang_vel: th.Tensor, + null_space_identity: th.Tensor, ): # Compute the inverse mm_inv = th.linalg.inv(mm) @@ -554,8 +562,9 @@ def _compute_osc_torques( # roboticsproceedings.org/rss07/p31.pdf if rest_qpos is not None: j_eef_inv = m_eef @ j_eef @ mm_inv + q_vec = q u_null = kd_null * -qd + kp_null * ((rest_qpos - q + math.pi) % (2 * math.pi) - math.pi) u_null = mm @ th.unsqueeze(u_null, dim=-1) - u += (th.eye(control_dim, dtype=th.float32) - j_eef.T @ j_eef_inv) @ u_null + u += (null_space_identity - j_eef.T @ j_eef_inv) @ u_null return u diff --git a/omnigibson/examples/action_primitives/rs_int_example.py b/omnigibson/examples/action_primitives/rs_int_example.py index fbecc2595..a0a99acbe 100644 --- a/omnigibson/examples/action_primitives/rs_int_example.py +++ b/omnigibson/examples/action_primitives/rs_int_example.py @@ -7,11 +7,6 @@ StarterSemanticActionPrimitives, StarterSemanticActionPrimitiveSet, ) -from omnigibson.macros import gm - -# Don't use GPU dynamics and use flatcache for performance boost -# gm.USE_GPU_DYNAMICS = True -# gm.ENABLE_FLATCACHE = True def execute_controller(ctrl_gen, env): diff --git a/omnigibson/examples/action_primitives/solve_simple_task.py b/omnigibson/examples/action_primitives/solve_simple_task.py index 168e3d239..4f72ae989 100644 --- a/omnigibson/examples/action_primitives/solve_simple_task.py +++ b/omnigibson/examples/action_primitives/solve_simple_task.py @@ -7,11 +7,6 @@ StarterSemanticActionPrimitives, StarterSemanticActionPrimitiveSet, ) -from omnigibson.macros import gm - -# Don't use GPU dynamics and use flatcache for performance boost -# gm.USE_GPU_DYNAMICS = True -# gm.ENABLE_FLATCACHE = True def execute_controller(ctrl_gen, env): diff --git a/omnigibson/examples/action_primitives/wip_solve_behavior_task.py b/omnigibson/examples/action_primitives/wip_solve_behavior_task.py index 3036fc42e..30c5334da 100644 --- a/omnigibson/examples/action_primitives/wip_solve_behavior_task.py +++ b/omnigibson/examples/action_primitives/wip_solve_behavior_task.py @@ -10,10 +10,6 @@ ) from omnigibson.macros import gm -# Don't use GPU dynamics and use flatcache for performance boost -# gm.USE_GPU_DYNAMICS = True -# gm.ENABLE_FLATCACHE = True - def execute_controller(ctrl_gen, env): for action in ctrl_gen: diff --git a/omnigibson/examples/environments/behavior_env_demo.py b/omnigibson/examples/environments/behavior_env_demo.py index 1fc460d83..e3933934e 100644 --- a/omnigibson/examples/environments/behavior_env_demo.py +++ b/omnigibson/examples/environments/behavior_env_demo.py @@ -1,5 +1,6 @@ import os +import torch as th import yaml import omnigibson as og @@ -8,7 +9,6 @@ # Make sure object states are enabled gm.ENABLE_OBJECT_STATES = True -gm.USE_GPU_DYNAMICS = True def main(random_selection=False, headless=False, short_exec=False): @@ -46,7 +46,7 @@ def main(random_selection=False, headless=False, short_exec=False): og.log.info("Resetting environment") env.reset() for i in range(100): - action = env.action_space.sample() + action = th.tensor(env.action_space.sample(), dtype=th.float32) state, reward, terminated, truncated, info = env.step(action) if terminated or truncated: og.log.info("Episode finished after {} timesteps".format(i + 1)) diff --git a/omnigibson/examples/environments/navigation_env_demo.py b/omnigibson/examples/environments/navigation_env_demo.py index e24be50da..3b7e711d3 100644 --- a/omnigibson/examples/environments/navigation_env_demo.py +++ b/omnigibson/examples/environments/navigation_env_demo.py @@ -1,5 +1,6 @@ import os +import torch as th import yaml import omnigibson as og @@ -41,7 +42,7 @@ def main(random_selection=False, headless=False, short_exec=False): og.log.info("Resetting environment") env.reset() for i in range(100): - action = env.action_space.sample() + action = th.tensor(env.action_space.sample(), dtype=th.float32) state, reward, terminated, truncated, info = env.step(action) if terminated or truncated: og.log.info("Episode finished after {} timesteps".format(i + 1)) diff --git a/omnigibson/examples/environments/vector_env_demo.py b/omnigibson/examples/environments/vector_env_demo.py index d0882e7a6..9dfb60764 100644 --- a/omnigibson/examples/environments/vector_env_demo.py +++ b/omnigibson/examples/environments/vector_env_demo.py @@ -17,8 +17,6 @@ def main(random_selection=False, headless=False, short_exec=False): # Load the config gm.RENDER_VIEWER_CAMERA = False - gm.ENABLE_FLATCACHE = True - gm.USE_GPU_DYNAMICS = False config_filename = os.path.join(og.example_config_path, "fetch_primitives.yaml") config = yaml.load(open(config_filename, "r"), Loader=yaml.FullLoader) diff --git a/omnigibson/examples/learning/navigation_policy_demo.py b/omnigibson/examples/learning/navigation_policy_demo.py index 33d905aa5..6c6651dbc 100644 --- a/omnigibson/examples/learning/navigation_policy_demo.py +++ b/omnigibson/examples/learning/navigation_policy_demo.py @@ -40,10 +40,9 @@ assert meets_minimum_version(gym.__version__, "0.28.1"), "Please install/update gymnasium to version >= 0.28.1" -# We don't need object states nor transitions rules, so we disable them now, and also enable flatcache for maximum speed +# We don't need object states nor transitions rules, so we disable them now gm.ENABLE_OBJECT_STATES = False gm.ENABLE_TRANSITION_RULES = False -gm.ENABLE_FLATCACHE = True class CustomCombinedExtractor(BaseFeaturesExtractor): @@ -170,7 +169,7 @@ def main(): policy_kwargs=policy_kwargs, n_steps=20 * 10, batch_size=8, - device="cuda", + device=og.sim.device, ) checkpoint_callback = CheckpointCallback(save_freq=1000, save_path=tensorboard_log_dir, name_prefix=prefix) eval_callback = EvalCallback(eval_env=env, eval_freq=1000, n_eval_episodes=20) diff --git a/omnigibson/examples/object_states/dicing_demo.py b/omnigibson/examples/object_states/dicing_demo.py index 6141e8fee..0f5012499 100644 --- a/omnigibson/examples/object_states/dicing_demo.py +++ b/omnigibson/examples/object_states/dicing_demo.py @@ -6,9 +6,8 @@ import omnigibson.utils.transform_utils as T from omnigibson.macros import gm -# Make sure object states, GPU dynamics, and transition rules are enabled +# Make sure object states and transition rules are enabled gm.ENABLE_OBJECT_STATES = True -gm.USE_GPU_DYNAMICS = True gm.ENABLE_TRANSITION_RULES = True diff --git a/omnigibson/examples/object_states/folded_unfolded_state_demo.py b/omnigibson/examples/object_states/folded_unfolded_state_demo.py index 9715e3813..c5f527758 100644 --- a/omnigibson/examples/object_states/folded_unfolded_state_demo.py +++ b/omnigibson/examples/object_states/folded_unfolded_state_demo.py @@ -8,7 +8,6 @@ # Make sure object states and GPU dynamics are enabled (GPU dynamics needed for cloth) gm.ENABLE_OBJECT_STATES = True -gm.USE_GPU_DYNAMICS = True def main(random_selection=False, headless=False, short_exec=False): @@ -19,6 +18,7 @@ def main(random_selection=False, headless=False, short_exec=False): # Create the scene config to load -- empty scene + custom cloth object cfg = { + "env": {"device": "cuda:0"}, "scene": { "type": "Scene", }, diff --git a/omnigibson/examples/object_states/object_state_texture_demo.py b/omnigibson/examples/object_states/object_state_texture_demo.py index bd3e46caf..8c7f5368a 100644 --- a/omnigibson/examples/object_states/object_state_texture_demo.py +++ b/omnigibson/examples/object_states/object_state_texture_demo.py @@ -5,9 +5,8 @@ from omnigibson.macros import gm, macros from omnigibson.utils.constants import ParticleModifyMethod -# Make sure object states are enabled, we're using GPU dynamics, and HQ rendering is enabled +# Make sure object states are enabled, and HQ rendering is enabled gm.ENABLE_OBJECT_STATES = True -gm.USE_GPU_DYNAMICS = True gm.ENABLE_HQ_RENDERING = True @@ -43,7 +42,7 @@ def main(random_selection=False, headless=False, short_exec=False): # In this case, we only allow our cabinet to absorb water, with no conditions needed. # This is needed for the Saturated ("saturable") state so that we can modify the texture # according to the water. - # NOTE: This will only change color if gm.ENABLE_HQ_RENDERING and gm.USE_GPU_DYNAMICS is + # NOTE: This will only change color if gm.ENABLE_HQ_RENDERING is # enabled! "water": [], }, diff --git a/omnigibson/examples/object_states/overlaid_demo.py b/omnigibson/examples/object_states/overlaid_demo.py index 9f6b1e334..9eee5a6e1 100644 --- a/omnigibson/examples/object_states/overlaid_demo.py +++ b/omnigibson/examples/object_states/overlaid_demo.py @@ -5,9 +5,7 @@ from omnigibson.object_states import Overlaid from omnigibson.utils.constants import PrimType -# Make sure object states and GPU dynamics are enabled (GPU dynamics needed for cloth) gm.ENABLE_OBJECT_STATES = True -gm.USE_GPU_DYNAMICS = True def main(random_selection=False, headless=False, short_exec=False): @@ -21,6 +19,9 @@ def main(random_selection=False, headless=False, short_exec=False): # Create the scene config to load -- empty scene + custom cloth object + custom rigid object cfg = { + "env": { + "device": "cuda", + }, "scene": { "type": "Scene", }, @@ -66,7 +67,10 @@ def main(random_selection=False, headless=False, short_exec=False): print("\nTry dragging cloth around with CTRL + Left-Click to see the Overlaid state change:\n") while steps != max_steps: - print(f"Overlaid {carpet.states[Overlaid].get_value(breakfast_table)} ", end="\r") + print( + f"Overlaid {carpet.states[Overlaid].get_value(breakfast_table)}. Avg keypoint pos {th.mean(carpet.root_link.keypoint_particle_positions, axis=0)} ", + end="\r", + ) env.step(th.empty(0)) steps += 1 diff --git a/omnigibson/examples/object_states/particle_applier_remover_demo.py b/omnigibson/examples/object_states/particle_applier_remover_demo.py index 24fbee551..fd05fec39 100644 --- a/omnigibson/examples/object_states/particle_applier_remover_demo.py +++ b/omnigibson/examples/object_states/particle_applier_remover_demo.py @@ -15,7 +15,6 @@ # Make sure object states and GPU dynamics are enabled (GPU dynamics needed for fluids) gm.ENABLE_OBJECT_STATES = True -gm.USE_GPU_DYNAMICS = True gm.ENABLE_HQ_RENDERING = True diff --git a/omnigibson/examples/object_states/particle_source_sink_demo.py b/omnigibson/examples/object_states/particle_source_sink_demo.py index 4add8f9bf..6511d6ce1 100644 --- a/omnigibson/examples/object_states/particle_source_sink_demo.py +++ b/omnigibson/examples/object_states/particle_source_sink_demo.py @@ -7,7 +7,6 @@ # Make sure object states are enabled and GPU dynamics are used gm.ENABLE_OBJECT_STATES = True -gm.USE_GPU_DYNAMICS = True gm.ENABLE_HQ_RENDERING = True diff --git a/omnigibson/examples/robots/grasping_mode_example.py b/omnigibson/examples/robots/grasping_mode_example.py index 89b5baf37..2341e0278 100644 --- a/omnigibson/examples/robots/grasping_mode_example.py +++ b/omnigibson/examples/robots/grasping_mode_example.py @@ -15,10 +15,6 @@ physical="Physical Grasping - No additional grasping assistance applied", ) -# Don't use GPU dynamics and Use flatcache for performance boost -gm.USE_GPU_DYNAMICS = False -gm.ENABLE_FLATCACHE = True - def main(random_selection=False, headless=False, short_exec=False): """ diff --git a/omnigibson/examples/robots/robot_control_example.py b/omnigibson/examples/robots/robot_control_example.py index 92a664e81..6cd76beff 100644 --- a/omnigibson/examples/robots/robot_control_example.py +++ b/omnigibson/examples/robots/robot_control_example.py @@ -22,10 +22,6 @@ empty="Empty environment with no objects", ) -# Don't use GPU dynamics and use flatcache for performance boost -gm.USE_GPU_DYNAMICS = False -gm.ENABLE_FLATCACHE = True - def choose_controllers(robot, random_selection=False): """ diff --git a/omnigibson/examples/scenes/scene_selector.py b/omnigibson/examples/scenes/scene_selector.py index 6ac5d52d1..c17874dd2 100644 --- a/omnigibson/examples/scenes/scene_selector.py +++ b/omnigibson/examples/scenes/scene_selector.py @@ -1,11 +1,11 @@ +import torch as th + import omnigibson as og from omnigibson.macros import gm from omnigibson.utils.asset_utils import get_available_g_scenes, get_available_og_scenes from omnigibson.utils.ui_utils import choose_from_options # Configure macros for maximum performance -gm.USE_GPU_DYNAMICS = True -gm.ENABLE_FLATCACHE = True gm.ENABLE_OBJECT_STATES = False gm.ENABLE_TRANSITION_RULES = False @@ -67,7 +67,7 @@ def main(random_selection=False, headless=False, short_exec=False): og.log.info("Resetting environment") env.reset() for i in range(100): - action = env.action_space.sample() + action = th.tensor(env.action_space.sample(), dtype=th.float32) state, reward, terminated, truncated, info = env.step(action) if terminated or truncated: og.log.info("Episode finished after {} timesteps".format(i + 1)) diff --git a/omnigibson/macros.py b/omnigibson/macros.py index b9f730808..1bb3e79d9 100644 --- a/omnigibson/macros.py +++ b/omnigibson/macros.py @@ -91,16 +91,9 @@ def determine_gm_path(default_path, env_var_name): # Whether to print out disclaimers (i.e.: known failure cases resulting from Omniverse's current bugs / limitations) gm.SHOW_DISCLAIMERS = False -# Whether to use omni's GPU dynamics -# This is necessary for certain features; e.g. particles (fluids / cloth) -gm.USE_GPU_DYNAMICS = False - # Whether to use high-fidelity rendering (this includes, e.g., isosurfaces) gm.ENABLE_HQ_RENDERING = False -# Whether to use omni's flatcache feature or not (can speed up simulation) -gm.ENABLE_FLATCACHE = False - # Whether to use continuous collision detection or not (slower simulation, but can prevent # objects from tunneling through each other) gm.ENABLE_CCD = False @@ -120,10 +113,10 @@ def determine_gm_path(default_path, env_var_name): gm.GPU_MAX_RIGID_PATCH_COUNT = 81920 * 4 # Whether to enable object state logic or not -gm.ENABLE_OBJECT_STATES = True +gm.ENABLE_OBJECT_STATES = False # Whether to enable transition rules or not -gm.ENABLE_TRANSITION_RULES = True +gm.ENABLE_TRANSITION_RULES = False # Default settings for the omni UI viewer gm.DEFAULT_VIEWER_WIDTH = 1280 diff --git a/omnigibson/object_states/__init__.py b/omnigibson/object_states/__init__.py index c6d487333..bf07236c2 100644 --- a/omnigibson/object_states/__init__.py +++ b/omnigibson/object_states/__init__.py @@ -22,7 +22,6 @@ from omnigibson.object_states.on_top import OnTop from omnigibson.object_states.open_state import Open from omnigibson.object_states.overlaid import Overlaid -from omnigibson.object_states.particle import ParticleRequirement from omnigibson.object_states.particle_modifier import ParticleApplier, ParticleRemover from omnigibson.object_states.particle_source_or_sink import ParticleSink, ParticleSource from omnigibson.object_states.pose import Pose diff --git a/omnigibson/object_states/cloth_mixin.py b/omnigibson/object_states/cloth_mixin.py index 3d947e5ca..cf0988d67 100644 --- a/omnigibson/object_states/cloth_mixin.py +++ b/omnigibson/object_states/cloth_mixin.py @@ -1,3 +1,4 @@ +import omnigibson as og from omnigibson.macros import gm from omnigibson.object_states.object_state_base import BaseObjectState from omnigibson.utils.constants import PrimType @@ -24,9 +25,9 @@ def is_compatible(cls, obj, **kwargs): f"with prim_type=PrimType.CLOTH!", ) - # Check for GPU dynamics - if not gm.USE_GPU_DYNAMICS: - return False, f"gm.USE_GPU_DYNAMICS must be True in order to use object state {cls.__name__}." + # Check for GPU pipeline + if og.sim.device == "cpu": + return False, f"Must be using GPU dynamics in order to use object state {cls.__name__}." return True, None diff --git a/omnigibson/object_states/factory.py b/omnigibson/object_states/factory.py index c4d26c88c..fd21f6c34 100644 --- a/omnigibson/object_states/factory.py +++ b/omnigibson/object_states/factory.py @@ -15,10 +15,10 @@ _ABILITY_DEPENDENCIES = { "robot": AbilityDependencies(states=[IsGrasping, ObjectsInFOVOfRobot], requirements=[]), "attachable": AbilityDependencies(states=[AttachedTo], requirements=[]), - "particleApplier": AbilityDependencies(states=[ParticleApplier], requirements=[ParticleRequirement]), - "particleRemover": AbilityDependencies(states=[ParticleRemover], requirements=[ParticleRequirement]), - "particleSource": AbilityDependencies(states=[ParticleSource], requirements=[ParticleRequirement]), - "particleSink": AbilityDependencies(states=[ParticleSink], requirements=[ParticleRequirement]), + "particleApplier": AbilityDependencies(states=[ParticleApplier], requirements=[]), + "particleRemover": AbilityDependencies(states=[ParticleRemover], requirements=[]), + "particleSource": AbilityDependencies(states=[ParticleSource], requirements=[]), + "particleSink": AbilityDependencies(states=[ParticleSink], requirements=[]), "coldSource": AbilityDependencies(states=[HeatSourceOrSink], requirements=[]), "cookable": AbilityDependencies(states=[Cooked, Burnt], requirements=[]), "coverable": AbilityDependencies(states=[Covered], requirements=[]), diff --git a/omnigibson/object_states/folded.py b/omnigibson/object_states/folded.py index b316b3441..382393b6a 100644 --- a/omnigibson/object_states/folded.py +++ b/omnigibson/object_states/folded.py @@ -61,7 +61,7 @@ def calculate_smoothness(self): normals = cloth.compute_face_normals(face_ids=cloth.keyface_idx) # projection onto the z-axis - proj = th.abs(normals @ th.tensor([0.0, 0.0, 1.0], dtype=th.float32)) + proj = th.abs(normals @ th.tensor([0.0, 0.0, 1.0], dtype=th.float32, device=normals.device)) percentage = th.mean((proj > math.cos(m.NORMAL_Z_ANGLE_DIFF)).float()).item() return percentage @@ -100,7 +100,7 @@ def calculate_projection_area_and_diagonal(self, dims): diagonal (float): diagonal of the convex hull of the projected points """ cloth = self.obj.root_link - points = cloth.keypoint_particle_positions[:, dims] + points = cloth.keypoint_particle_positions[:, dims].cpu() try: hull = ConvexHull(points) diff --git a/omnigibson/object_states/overlaid.py b/omnigibson/object_states/overlaid.py index 12f6369c5..20b73d08b 100644 --- a/omnigibson/object_states/overlaid.py +++ b/omnigibson/object_states/overlaid.py @@ -56,14 +56,14 @@ def _get_value(self, other): return False # Compute the convex hull of the particles of the cloth object. - points = self.obj.root_link.keypoint_particle_positions[:, :2] + points = self.obj.root_link.keypoint_particle_positions[:, :2].cpu() cloth_hull = ConvexHull(points) # Compute the base aligned bounding box of the rigid object. bbox_center, bbox_orn, bbox_extent, _ = other.get_base_aligned_bbox(xy_aligned=True) vertices_local = th.tensor(list(itertools.product((1, -1), repeat=3))) * (bbox_extent / 2) vertices = T.transform_points(vertices_local, T.pose2mat((bbox_center, bbox_orn))) - rigid_hull = ConvexHull(vertices[:, :2]) + rigid_hull = ConvexHull(vertices[:, :2].cpu()) # The goal is to find the intersection of the convex hull and the bounding box. # We can do so with HalfspaceIntersection, which takes as input a list of equations that define the half spaces, @@ -71,7 +71,7 @@ def _get_value(self, other): interior_pt = th.mean(vertices, dim=0)[:2] half_spaces = th.vstack((th.tensor(cloth_hull.equations), th.tensor(rigid_hull.equations))) try: - half_space_intersection = HalfspaceIntersection(half_spaces, interior_pt) + half_space_intersection = HalfspaceIntersection(half_spaces.cpu(), interior_pt.cpu()) except QhullError: # The bbox center of the rigid body does not lie in the intersection, return False. return False diff --git a/omnigibson/object_states/particle.py b/omnigibson/object_states/particle.py deleted file mode 100644 index 25c2ea4df..000000000 --- a/omnigibson/object_states/particle.py +++ /dev/null @@ -1,18 +0,0 @@ -import torch as th - -from omnigibson.object_states.object_state_base import BaseObjectRequirement - - -class ParticleRequirement(BaseObjectRequirement): - """ - Class for sanity checking objects that requires particle systems - """ - - @classmethod - def is_compatible(cls, obj, **kwargs): - from omnigibson.macros import gm - - if not gm.USE_GPU_DYNAMICS: - return False, f"Particle systems are not enabled when GPU dynamics is off." - - return True, None diff --git a/omnigibson/object_states/particle_modifier.py b/omnigibson/object_states/particle_modifier.py index a7f655eab..610dbfb1e 100644 --- a/omnigibson/object_states/particle_modifier.py +++ b/omnigibson/object_states/particle_modifier.py @@ -642,15 +642,10 @@ def check_conditions_for_system(self, system_name): return all(condition(self.obj) for condition in self.conditions[system_name]) def _update(self): - # If we're using projection method and flatcache, we need to manually update this object's transforms on the USD + # If we're using projection method, we need to manually update this object's transforms on the USD # so the corresponding visualization and overlap meshes are updated properly # This is expensive, so only do it if the object is not a fixed object and we have an active projection - if ( - self.method == ParticleModifyMethod.PROJECTION - and gm.ENABLE_FLATCACHE - and not self.obj.fixed_base - and self.projection_is_active - ): + if self.method == ParticleModifyMethod.PROJECTION and not self.obj.fixed_base and self.projection_is_active: FlatcacheAPI.sync_raw_object_transforms_in_usd(prim=self.obj) # Check if there's any overlap and if we're at the correct step @@ -1490,7 +1485,7 @@ def systems_to_check(self): @property def projection_is_active(self): # Only active if the projection mesh is enabled - return self.projection_emitter.GetProperty("inputs:active").Get() + return self.projection_emitter.GetProperty("inputs:active").Get() if self.visualize else False @classproperty def metalink_prefix(cls): @@ -1501,19 +1496,6 @@ def requires_metalink(cls, **kwargs): # No metalink required for adjacency return kwargs.get("method", ParticleModifyMethod.ADJACENCY) != ParticleModifyMethod.ADJACENCY - @classmethod - def is_compatible(cls, obj, **kwargs): - # Run super first - compatible, reason = super().is_compatible(obj, **kwargs) - if not compatible: - return compatible, reason - - # Check whether GPU dynamics are enabled (necessary for this object state) - if not gm.USE_GPU_DYNAMICS: - return False, f"gm.USE_GPU_DYNAMICS must be True in order to use object state {cls.__name__}." - - return True, None - @property def _default_link(self): # Only supported for adjacency, NOT projection diff --git a/omnigibson/objects/controllable_object.py b/omnigibson/objects/controllable_object.py index 8feb77008..2bd928dea 100644 --- a/omnigibson/objects/controllable_object.py +++ b/omnigibson/objects/controllable_object.py @@ -270,7 +270,7 @@ def _load_controllers(self): if name in controller_subsumes: for subsumed_name in controller_subsumes[name]: subsumed_cfg = _controller_config[subsumed_name] - cfg["dof_idx"] = th.concatenate([subsumed_cfg["dof_idx"], cfg["dof_idx"]]) + cfg["dof_idx"] = subsumed_cfg["dof_idx"] + cfg["dof_idx"] # If we're using normalized action space, override the inputs for all controllers if self._action_normalize: @@ -295,8 +295,8 @@ def update_controller_mode(self): for name in self._controllers: for dof in self._controllers[name].dof_idx: # Make sure the DOF has not already been set yet, and remove it afterwards - assert dof.item() in unused_dofs - unused_dofs.remove(dof.item()) + assert dof in unused_dofs + unused_dofs.remove(dof) control_type = self._controllers[name].control_type self._joints[self.dof_names_ordered[dof]].set_control_type( control_type=control_type, @@ -404,13 +404,13 @@ def _create_continuous_action_space(self): low, high = [], [] for controller in self._controllers.values(): limits = controller.command_input_limits - low.append(th.tensor([-float("inf")] * controller.command_dim) if limits is None else limits[0]) - high.append(th.tensor([float("inf")] * controller.command_dim) if limits is None else limits[1]) + low.append(th.tensor([-float("inf")] * controller.command_dim) if limits is None else limits[0].cpu()) + high.append(th.tensor([float("inf")] * controller.command_dim) if limits is None else limits[1].cpu()) return gym.spaces.Box( shape=(self.action_dim,), - low=th.cat(low).cpu().numpy(), - high=th.cat(high).cpu().numpy(), + low=th.cat(low).numpy(), + high=th.cat(high).numpy(), dtype=NumpyTypes.FLOAT32, ) @@ -488,13 +488,12 @@ def step(self): idx += controller.command_dim # Compose controls - u_vec = th.zeros(self.n_dof) + u_vec = th.zeros(self.n_dof, device=og.sim.device) # By default, the control type is None and the control value is 0 (th.zeros) - i.e. no control applied u_type_vec = th.tensor([ControlType.NONE] * self.n_dof) for group, ctrl in control.items(): - idx = self._controllers[group].dof_idx - u_vec[idx] = ctrl["value"] - u_type_vec[idx] = ctrl["type"] + u_vec[self._controllers[group].dof_idx] = ctrl["value"] + u_type_vec[self._controllers[group].dof_idx] = ctrl["type"] u_vec, u_type_vec = self._postprocess_control(control=u_vec, control_type=u_type_vec) @@ -610,7 +609,7 @@ def deploy_control(self, control, control_type): using_pos = True elif ctrl_type == ControlType.NONE: # Set zero efforts - eff_vec.append(0) + eff_vec.append(th.tensor(0, dtype=th.float32, device=og.sim.device)) eff_idxs.append(cur_ctrl_idx) using_eff = True else: @@ -621,15 +620,15 @@ def deploy_control(self, control, control_type): # set the targets for joints if using_pos: ControllableObjectViewAPI.set_joint_position_targets( - self.articulation_root_path, positions=th.tensor(pos_vec, dtype=th.float), indices=th.tensor(pos_idxs) + self.articulation_root_path, positions=th.stack(pos_vec), indices=pos_idxs ) if using_vel: ControllableObjectViewAPI.set_joint_velocity_targets( - self.articulation_root_path, velocities=th.tensor(vel_vec, dtype=th.float), indices=th.tensor(vel_idxs) + self.articulation_root_path, velocities=th.stack(vel_vec), indices=vel_idxs ) if using_eff: ControllableObjectViewAPI.set_joint_efforts( - self.articulation_root_path, efforts=th.tensor(eff_vec, dtype=th.float), indices=th.tensor(eff_idxs) + self.articulation_root_path, efforts=th.stack(eff_vec), indices=eff_idxs ) def get_control_dict(self): diff --git a/omnigibson/objects/object_base.py b/omnigibson/objects/object_base.py index c34231551..a4101b42f 100644 --- a/omnigibson/objects/object_base.py +++ b/omnigibson/objects/object_base.py @@ -234,7 +234,7 @@ def _initialize(self): "emissive_intensity": material.emissive_intensity, } - @property + @cached_property def articulation_root_path(self): has_articulated_joints, has_fixed_joints = self.n_joints > 0, self.n_fixed_joints > 0 if self.kinematic_only or ((not has_articulated_joints) and (not has_fixed_joints)): diff --git a/omnigibson/prims/cloth_prim.py b/omnigibson/prims/cloth_prim.py index 0177d8fab..ba604aec8 100644 --- a/omnigibson/prims/cloth_prim.py +++ b/omnigibson/prims/cloth_prim.py @@ -17,10 +17,15 @@ import omnigibson.utils.transform_utils as T from omnigibson.macros import create_module_macros, gm from omnigibson.prims.geom_prim import GeomPrim +from omnigibson.utils.geometry_utils import get_particle_positions_from_frame, get_particle_positions_in_frame from omnigibson.utils.numpy_utils import vtarray_to_torch from omnigibson.utils.sim_utils import CsRawData +from omnigibson.utils.ui_utils import create_module_logger from omnigibson.utils.usd_utils import array_to_vtarray, mesh_prim_to_trimesh_mesh, sample_mesh_keypoints +# Create module logger +log = create_module_logger(module_name=__name__) + # Create settings for this module m = create_module_macros(module_path=__file__) @@ -62,6 +67,8 @@ def __init__( self._centroid_idx = None self._keypoint_idx = None self._keyface_idx = None + self._cloth_prim_view = None + self._cloth_view_initialized = False # Run super init super().__init__( @@ -74,10 +81,7 @@ def _post_load(self): # run super first super()._post_load() - # Make sure flatcache is not being used -- if so, raise an error, since we lose most of our needed functionality - # (such as R/W to specific particle states) when flatcache is enabled - assert not gm.ENABLE_FLATCACHE, "Cannot use flatcache with ClothPrim!" - + # Can we do this also using the Cloth API, in initialize? self._mass_api = ( lazy.pxr.UsdPhysics.MassAPI(self._prim) if self._prim.HasAPI(lazy.pxr.UsdPhysics.MassAPI) @@ -91,47 +95,50 @@ def _post_load(self): # Clothify this prim, which is assumed to be a mesh self.cloth_system.clothify_mesh_prim(mesh_prim=self._prim, remesh=self._load_config.get("remesh", True)) - # Track generated particle count - positions = self.compute_particle_positions() - self._n_particles = len(positions) - - # Sample mesh keypoints / keyvalues and sanity check the AABB of these subsampled points vs. the actual points - success = False - for i in range(10): - self._keypoint_idx, self._keyface_idx = sample_mesh_keypoints( - mesh_prim=self._prim, - n_keypoints=m.N_CLOTH_KEYPOINTS, - n_keyfaces=m.N_CLOTH_KEYFACES, - seed=i, - ) - - keypoint_positions = positions[self._keypoint_idx] - keypoint_aabb = keypoint_positions.min(dim=0).values, keypoint_positions.max(dim=0).values - true_aabb = positions.min(dim=0).values, positions.max(dim=0).values - overlap_x = th.max( - th.min(true_aabb[1][0], keypoint_aabb[1][0]) - th.max(true_aabb[0][0], keypoint_aabb[0][0]), - th.tensor(0), - ) - overlap_y = th.max( - th.min(true_aabb[1][1], keypoint_aabb[1][1]) - th.max(true_aabb[0][1], keypoint_aabb[0][1]), - th.tensor(0), - ) - overlap_z = th.max( - th.min(true_aabb[1][2], keypoint_aabb[1][2]) - th.max(true_aabb[0][2], keypoint_aabb[0][2]), - th.tensor(0), - ) - overlap_vol = overlap_x * overlap_y * overlap_z - true_vol = th.prod(true_aabb[1] - true_aabb[0]) - if true_vol == 0.0 or (overlap_vol / true_vol > m.KEYPOINT_COVERAGE_THRESHOLD).item(): - success = True - break - assert success, f"Did not adequately subsample keypoints for cloth {self.name}!" - - # Compute centroid particle idx based on AABB - aabb_min, aabb_max = th.min(positions, dim=0).values, th.max(positions, dim=0).values - aabb_center = (aabb_min + aabb_max) / 2.0 - dists = th.norm(positions - aabb_center.reshape(1, 3), dim=-1) - self._centroid_idx = th.argmin(dists) + # Track generated particle count. This is the only time we use the USD API. + self._n_particles = len(self._prim.GetAttribute("points").Get()) + + # Load the cloth prim view if using gpu pipeline + if og.sim.device.startswith("cuda"): + self._cloth_prim_view = lazy.omni.isaac.core.prims.ClothPrimView(self.prim_path) + + # # Sample mesh keypoints / keyvalues and sanity check the AABB of these subsampled points vs. the actual points + # success = False + # for i in range(10): + # self._keypoint_idx, self._keyface_idx = sample_mesh_keypoints( + # mesh_prim=self._prim, + # n_keypoints=m.N_CLOTH_KEYPOINTS, + # n_keyfaces=m.N_CLOTH_KEYFACES, + # seed=i, + # ) + + # keypoint_positions = positions[self._keypoint_idx] + # keypoint_aabb = keypoint_positions.min(dim=0).values, keypoint_positions.max(dim=0).values + # true_aabb = positions.min(dim=0).values, positions.max(dim=0).values + # overlap_x = th.max( + # th.min(true_aabb[1][0], keypoint_aabb[1][0]) - th.max(true_aabb[0][0], keypoint_aabb[0][0]), + # th.tensor(0), + # ) + # overlap_y = th.max( + # th.min(true_aabb[1][1], keypoint_aabb[1][1]) - th.max(true_aabb[0][1], keypoint_aabb[0][1]), + # th.tensor(0), + # ) + # overlap_z = th.max( + # th.min(true_aabb[1][2], keypoint_aabb[1][2]) - th.max(true_aabb[0][2], keypoint_aabb[0][2]), + # th.tensor(0), + # ) + # overlap_vol = overlap_x * overlap_y * overlap_z + # true_vol = th.prod(true_aabb[1] - true_aabb[0]) + # if true_vol == 0.0 or (overlap_vol / true_vol > m.KEYPOINT_COVERAGE_THRESHOLD).item(): + # success = True + # break + # assert success, f"Did not adequately subsample keypoints for cloth {self.name}!" + + # # Compute centroid particle idx based on AABB + # aabb_min, aabb_max = th.min(positions, dim=0).values, th.max(positions, dim=0).values + # aabb_center = (aabb_min + aabb_max) / 2.0 + # dists = th.norm(positions - aabb_center.reshape(1, 3), dim=-1) + # self._centroid_idx = th.argmin(dists) # Store the default position of the points in the local frame self._default_positions = vtarray_to_torch(self.get_attribute(attr="points")) @@ -168,6 +175,14 @@ def kinematic_only(self): """ return False + def _compute_usd_particle_positions(self, idxs=None): + # Don't copy to save compute, since we won't be returning a reference to the underlying object anyways + p_local = vtarray_to_torch(self.get_attribute(attr="points")) + + return get_particle_positions_in_frame( + *self.get_position_orientation(), self.scale, p_local[idxs] if idxs is not None else p_local + ) + def compute_particle_positions(self, idxs=None): """ Compute individual particle positions for this cloth prim @@ -179,16 +194,14 @@ def compute_particle_positions(self, idxs=None): th.tensor: (N, 3) numpy array, where each of the N particles' positions are expressed in (x,y,z) cartesian coordinates relative to the world frame """ - pos, ori = self.get_position_orientation() - ori = T.quat2mat(ori) - scale = self.scale - - # Don't copy to save compute, since we won't be returning a reference to the underlying object anyways - p_local = vtarray_to_torch(self.get_attribute(attr="points")) - p_local = p_local[idxs] if idxs is not None else p_local - p_world = (ori @ (p_local * scale).T).T + pos - - return p_world + if og.sim.device.startswith("cuda"): + all_particle_positions = self._cloth_prim_view.get_world_positions(clone=False)[0, :, :] + return all_particle_positions[: self._n_particles] if idxs is None else all_particle_positions[idxs] + else: + log.warning( + "CPU pipeline used for particle position computation, this is not supported. Cloth particle positions will not update." + ) + return self._compute_usd_particle_positions(idxs=idxs) def set_particle_positions(self, positions, idxs=None): """ @@ -199,23 +212,24 @@ def set_particle_positions(self, positions, idxs=None): cartesian coordinates relative to the world frame idxs (n-array or None): If set, will only set the requested indexed particle state """ - n_expected = self._n_particles if idxs is None else len(idxs) - assert ( - len(positions) == n_expected - ), f"Got mismatch in particle setting size: {len(positions)}, vs. number of expected particles {n_expected}!" + if og.sim.device.startswith("cuda"): + n_expected = self._n_particles if idxs is None else len(idxs) + assert ( + len(positions) == n_expected + ), f"Got mismatch in particle setting size: {len(positions)}, vs. number of expected particles {n_expected}!" - translation, rotation = self.get_position_orientation() - rotation = T.quat2mat(rotation) - scale = self.scale - p_local = (rotation.T @ (positions - translation).T).T / scale + # First, get the particle positions. + cur_pos = self._cloth_prim_view.get_world_positions() - # Fill the idxs if requested - if idxs is not None: - p_local_old = vtarray_to_torch(self.get_attribute(attr="points")) - p_local_old[idxs] = p_local - p_local = p_local_old + # Then apply the new positions at the appropriate indices + cur_pos[0, idxs] = positions - self.set_attribute(attr="points", val=lazy.pxr.Vt.Vec3fArray(p_local.tolist())) + # Then set to that position + self._cloth_prim_view.set_world_positions(cur_pos) + else: + log.warning( + "CPU pipeline used for particle position setting, this is not supported. Cloth particle positions will not update." + ) @property def keypoint_idx(self): @@ -372,8 +386,14 @@ def report_hit(hit): return contacts def update_handles(self): - # no handles to update - pass + if og.sim.device.startswith("cuda"): + assert og.sim._physics_sim_view._backend is not None, "Physics sim backend not initialized!" + self._cloth_prim_view.initialize(og.sim.physics_sim_view) + self._cloth_view_initialized = True + assert self._n_particles <= self._cloth_prim_view.max_particles_per_cloth, ( + f"Got more particles than the maximum allowed for this cloth! Got {self._n_particles}, max is " + f"{self._cloth_prim_view.max_particles_per_cloth}!" + ) @property def volume(self): @@ -541,7 +561,11 @@ def _dump_state(self): state = super()._dump_state() state["particle_group"] = self.particle_group state["n_particles"] = self.n_particles - state["particle_positions"] = self.compute_particle_positions() + state["particle_positions"] = ( + self.compute_particle_positions().cpu() + if self._cloth_view_initialized + else self._compute_usd_particle_positions() + ) state["particle_velocities"] = self.particle_velocities return state @@ -612,5 +636,7 @@ def reset(self): Reset the points to their default positions in the local frame, and also zeroes out velocities """ if self.initialized: - self.set_attribute(attr="points", val=lazy.pxr.Vt.Vec3fArray(self._default_positions.tolist())) + self.set_particle_positions( + get_particle_positions_in_frame(*self.get_position_orientation(), self.scale, self._default_positions) + ) self.particle_velocities = th.zeros((self._n_particles, 3)) diff --git a/omnigibson/prims/entity_prim.py b/omnigibson/prims/entity_prim.py index 6676097fc..c58074bda 100644 --- a/omnigibson/prims/entity_prim.py +++ b/omnigibson/prims/entity_prim.py @@ -885,7 +885,7 @@ def get_joint_positions(self, normalized=False): # Run sanity checks -- make sure we are articulated assert self.n_joints > 0, "Tried to call method not intended for entity prim with no joints!" - joint_positions = self._articulation_view.get_joint_positions().view(self.n_dof) + joint_positions = self._articulation_view.get_joint_positions().view(self.n_dof).cpu() # Possibly normalize values when returning return self._normalize_positions(positions=joint_positions) if normalized else joint_positions @@ -903,7 +903,7 @@ def get_joint_velocities(self, normalized=False): # Run sanity checks -- make sure we are articulated assert self.n_joints > 0, "Tried to call method not intended for entity prim with no joints!" - joint_velocities = self._articulation_view.get_joint_velocities().view(self.n_dof) + joint_velocities = self._articulation_view.get_joint_velocities().view(self.n_dof).cpu() # Possibly normalize values when returning return self._normalize_velocities(velocities=joint_velocities) if normalized else joint_velocities @@ -921,7 +921,7 @@ def get_joint_efforts(self, normalized=False): # Run sanity checks -- make sure we are articulated assert self.n_joints > 0, "Tried to call method not intended for entity prim with no joints!" - joint_efforts = self._articulation_view.get_measured_joint_efforts().view(self.n_dof) + joint_efforts = self._articulation_view.get_measured_joint_efforts().view(self.n_dof).cpu() # Possibly normalize values when returning return self._normalize_efforts(efforts=joint_efforts) if normalized else joint_efforts @@ -939,7 +939,7 @@ def get_joint_position_targets(self, normalized=False): # Run sanity checks -- make sure we are articulated assert self.n_joints > 0, "Tried to call method not intended for entity prim with no joints!" - joint_positions = self._articulation_view.get_joint_position_targets().view(self.n_dof) + joint_positions = self._articulation_view.get_joint_position_targets().view(self.n_dof).cpu() # Possibly normalize values when returning return self._normalize_positions(positions=joint_positions) if normalized else joint_positions @@ -957,7 +957,7 @@ def get_joint_velocity_targets(self, normalized=False): # Run sanity checks -- make sure we are articulated assert self.n_joints > 0, "Tried to call method not intended for entity prim with no joints!" - joint_velocities = self._articulation_view.get_joint_velocity_targets().view(self.n_dof) + joint_velocities = self._articulation_view.get_joint_velocity_targets().view(self.n_dof).cpu() # Possibly normalize values when returning return self._normalize_velocities(velocities=joint_velocities) if normalized else joint_velocities @@ -1091,8 +1091,8 @@ def get_position_orientation(self, frame: Literal["world", "scene"] = "world", c # Otherwise, get the pose from the articulation view and convert to our format positions, orientations = self._articulation_view.get_world_poses(clone=clone) - position = positions[0] - orientation = orientations[0][[1, 2, 3, 0]] + position = positions[0].cpu() + orientation = orientations[0][[1, 2, 3, 0]].cpu() # Assert that the orientation is a unit quaternion assert math.isclose( @@ -1642,8 +1642,8 @@ def serialize(self, state): state_flat = [self.root_link.serialize(state=state["root_link"])] if self.n_joints > 0: state_flat += [ - state["joint_pos"], - state["joint_vel"], + state["joint_pos"].to(device="cpu"), + state["joint_vel"].to(device="cpu"), ] return th.cat(state_flat) diff --git a/omnigibson/prims/geom_prim.py b/omnigibson/prims/geom_prim.py index 7b03dad9e..98741e419 100644 --- a/omnigibson/prims/geom_prim.py +++ b/omnigibson/prims/geom_prim.py @@ -260,8 +260,8 @@ def collision_enabled(self, enabled): """ # Currently, trying to toggle while simulator is playing while using GPU dynamics results in a crash, so we # assert that the sim is stopped here - if self._initialized and gm.USE_GPU_DYNAMICS: - assert og.sim.is_stopped(), "Cannot toggle collisions while using GPU dynamics unless simulator is stopped!" + if self._initialized: + assert og.sim.is_stopped(), "Cannot toggle collisions unless simulator is stopped!" self.set_attribute("physics:collisionEnabled", enabled) # TODO: Maybe this should all be added to RigidPrim instead? diff --git a/omnigibson/prims/rigid_prim.py b/omnigibson/prims/rigid_prim.py index 6167f4ea9..46dd09e56 100644 --- a/omnigibson/prims/rigid_prim.py +++ b/omnigibson/prims/rigid_prim.py @@ -275,14 +275,25 @@ def contact_list(self): contacts.append(CsRawData(*c)) return contacts - def set_linear_velocity(self, velocity): + def set_velocity(self, velocity): + """ + Sets the linear and angular velocity of the prim in stage. + Args: + velocity (th.Tensor): linear and angular velocity to set the rigid prim to. Shape (6,). + """ + assert velocity.shape == (6,), f"Velocity must be a 6-array, got {velocity.shape}" + self._rigid_prim_view.set_velocities(velocity[None, :]) + + def set_linear_velocity(self, linear_velocity): """ Sets the linear velocity of the prim in stage. Args: - velocity (th.tensor): linear velocity to set the rigid prim to. Shape (3,). + linear_velocity (th.tensor): linear velocity to set the rigid prim to. Shape (3,). """ - self._rigid_prim_view.set_linear_velocities(velocity[None, :]) + ang_vel = self.get_angular_velocity() + vel = th.cat([linear_velocity, ang_vel]) + self.set_velocity(vel) def get_linear_velocity(self, clone=True): """ @@ -292,16 +303,18 @@ def get_linear_velocity(self, clone=True): Returns: th.tensor: current linear velocity of the the rigid prim. Shape (3,). """ - return self._rigid_prim_view.get_linear_velocities(clone=clone)[0] + return self._rigid_prim_view.get_linear_velocities(clone=clone)[0].cpu() - def set_angular_velocity(self, velocity): + def set_angular_velocity(self, angular_velocity): """ Sets the angular velocity of the prim in stage. Args: velocity (th.tensor): angular velocity to set the rigid prim to. Shape (3,). """ - self._rigid_prim_view.set_angular_velocities(velocity[None, :]) + lin_vel = self.get_linear_velocity() + vel = th.cat([lin_vel, angular_velocity]) + self.set_velocity(vel) def get_angular_velocity(self, clone=True): """ @@ -311,7 +324,7 @@ def get_angular_velocity(self, clone=True): Returns: th.tensor: current angular velocity of the the rigid prim. Shape (3,). """ - return self._rigid_prim_view.get_angular_velocities(clone=clone)[0] + return self._rigid_prim_view.get_angular_velocities(clone=clone)[0].cpu() def set_position_orientation(self, position=None, orientation=None, frame: Literal["world", "scene"] = "world"): """ @@ -379,8 +392,8 @@ def get_position_orientation(self, frame: Literal["world", "scene"] = "world", c # Otherwise, get the pose from the rigid prim view and convert to our format positions, orientations = self._rigid_prim_view.get_world_poses(clone=clone) - position = positions[0] - orientation = orientations[0][[1, 2, 3, 0]] + position = positions[0].cpu() + orientation = orientations[0].cpu()[[1, 2, 3, 0]] # Assert that the orientation is a unit quaternion assert math.isclose( @@ -833,8 +846,8 @@ def clear_kinematic_only_cache(self): def _dump_state(self): # Grab pose from super class state = super()._dump_state() - state["lin_vel"] = self.get_linear_velocity(clone=False) - state["ang_vel"] = self.get_angular_velocity(clone=False) + state["lin_vel"] = self.get_linear_velocity(clone=False).to("cpu") + state["ang_vel"] = self.get_angular_velocity(clone=False).to("cpu") return state diff --git a/omnigibson/prims/xform_prim.py b/omnigibson/prims/xform_prim.py index 27f71f8f0..3e4018e51 100644 --- a/omnigibson/prims/xform_prim.py +++ b/omnigibson/prims/xform_prim.py @@ -229,17 +229,16 @@ def set_position_orientation( rotq = lazy.pxr.Gf.Quatd(*orientation) xform_op.Set(rotq) PoseAPI.invalidate() - if gm.ENABLE_FLATCACHE: - # If flatcache is on, make sure the USD local pose is synced to the fabric local pose. - # Ideally we should call usdrt's set local pose directly, but there is no such API. - # The only available API is SetLocalXformFromUsd, so we update USD first, and then sync to fabric. - xformable_prim = lazy.usdrt.Rt.Xformable( - lazy.omni.isaac.core.utils.prims.get_prim_at_path(self.prim_path, fabric=True) - ) - assert ( - not xformable_prim.HasWorldXform() - ), "Fabric's world pose is set for a non-rigid prim which is unexpected. Please report this." - xformable_prim.SetLocalXformFromUsd() + # If flatcache is on, make sure the USD local pose is synced to the fabric local pose. + # Ideally we should call usdrt's set local pose directly, but there is no such API. + # The only available API is SetLocalXformFromUsd, so we update USD first, and then sync to fabric. + xformable_prim = lazy.usdrt.Rt.Xformable( + lazy.omni.isaac.core.utils.prims.get_prim_at_path(self.prim_path, fabric=True) + ) + assert ( + not xformable_prim.HasWorldXform() + ), "Fabric's world pose is set for a non-rigid prim which is unexpected. Please report this." + xformable_prim.SetLocalXformFromUsd() def get_position_orientation(self, frame: Literal["world", "scene", "parent"] = "world", clone=True): """ diff --git a/omnigibson/robots/active_camera_robot.py b/omnigibson/robots/active_camera_robot.py index cb56e4c77..258ddf7ab 100644 --- a/omnigibson/robots/active_camera_robot.py +++ b/omnigibson/robots/active_camera_robot.py @@ -132,7 +132,7 @@ def camera_control_idx(self): Returns: n-array: Indices in low-level control vector corresponding to camera joints. """ - return th.tensor([list(self.joints.keys()).index(name) for name in self.camera_joint_names]) + return [list(self.joints.keys()).index(name) for name in self.camera_joint_names] @classproperty def _do_not_register_classes(cls): diff --git a/omnigibson/robots/articulated_trunk_robot.py b/omnigibson/robots/articulated_trunk_robot.py index 889d0e10d..adf9fd874 100644 --- a/omnigibson/robots/articulated_trunk_robot.py +++ b/omnigibson/robots/articulated_trunk_robot.py @@ -51,7 +51,7 @@ def trunk_control_idx(self): Returns: n-array: Indices in low-level control vector corresponding to trunk joints. """ - return th.tensor([list(self.joints.keys()).index(name) for name in self.trunk_joint_names]) + return [list(self.joints.keys()).index(name) for name in self.trunk_joint_names] @property def trunk_action_idx(self): diff --git a/omnigibson/robots/locomotion_robot.py b/omnigibson/robots/locomotion_robot.py index 95601c175..f28a49331 100644 --- a/omnigibson/robots/locomotion_robot.py +++ b/omnigibson/robots/locomotion_robot.py @@ -224,7 +224,7 @@ def base_control_idx(self): Returns: n-array: Indices in low-level control vector corresponding to base joints. """ - return th.tensor([list(self.joints.keys()).index(name) for name in self.base_joint_names]) + return [list(self.joints.keys()).index(name) for name in self.base_joint_names] @classproperty def _do_not_register_classes(cls): diff --git a/omnigibson/robots/manipulation_robot.py b/omnigibson/robots/manipulation_robot.py index f3f6d55ff..0234fc587 100644 --- a/omnigibson/robots/manipulation_robot.py +++ b/omnigibson/robots/manipulation_robot.py @@ -1,6 +1,7 @@ import math from abc import abstractmethod from collections import namedtuple +from functools import cached_property from typing import Literal import networkx as nx @@ -600,8 +601,7 @@ def arm_control_idx(self): vector corresponding to arm joints. """ return { - arm: th.tensor([list(self.joints.keys()).index(name) for name in self.arm_joint_names[arm]]) - for arm in self.arm_names + arm: [list(self.joints.keys()).index(name) for name in self.arm_joint_names[arm]] for arm in self.arm_names } @property @@ -612,7 +612,7 @@ def gripper_control_idx(self): vector corresponding to gripper joints. """ return { - arm: th.tensor([list(self.joints.keys()).index(name) for name in self.finger_joint_names[arm]]) + arm: [list(self.joints.keys()).index(name) for name in self.finger_joint_names[arm]] for arm in self.arm_names } @@ -625,7 +625,7 @@ def arm_links(self): """ return {arm: [self._links[link] for link in self.arm_link_names[arm]] for arm in self.arm_names} - @property + @cached_property def eef_links(self): """ Returns: diff --git a/omnigibson/robots/tiago.py b/omnigibson/robots/tiago.py index 43614b8b4..a146620cb 100644 --- a/omnigibson/robots/tiago.py +++ b/omnigibson/robots/tiago.py @@ -272,7 +272,7 @@ def arm_control_idx(self): # Add combined entry idxs = super().arm_control_idx # Concatenate all values and sort them - idxs["combined"] = th.sort(th.cat([val for val in idxs.values()]))[0] + idxs["combined"] = sorted([idx for arm, arm_idxes in idxs.items() for idx in arm_idxes]) return idxs @property diff --git a/omnigibson/scenes/scene_base.py b/omnigibson/scenes/scene_base.py index acb76ec18..614a129d1 100644 --- a/omnigibson/scenes/scene_base.py +++ b/omnigibson/scenes/scene_base.py @@ -348,10 +348,7 @@ def _load_scene_prim_with_objects(self, last_scene_edge, initial_scene_prim_z_of # Create desired systems for system_name in self._init_systems: - if gm.USE_GPU_DYNAMICS: - self.get_system(system_name) - else: - log.warning(f"System {system_name} is not supported without GPU dynamics! Skipping...") + self.get_system(system_name) # Position the scene prim initially at a z offset to avoid collision self._scene_prim.set_position_orientation( diff --git a/omnigibson/simulator.py b/omnigibson/simulator.py index ce06d7a21..84f7bb87d 100644 --- a/omnigibson/simulator.py +++ b/omnigibson/simulator.py @@ -449,15 +449,9 @@ def _set_physics_engine_settings(self): # default collide with each other, and modify settings for speed optimization self._physics_context.set_invert_collision_group_filter(False) self._physics_context.enable_ccd(gm.ENABLE_CCD) - self._physics_context.enable_fabric(gm.ENABLE_FLATCACHE) - - # Enable GPU dynamics based on whether we need omni particles feature - if gm.USE_GPU_DYNAMICS: - self._physics_context.enable_gpu_dynamics(True) - self._physics_context.set_broadphase_type("GPU") - else: - self._physics_context.enable_gpu_dynamics(False) - self._physics_context.set_broadphase_type("MBP") + self._physics_context.enable_fabric(True) + self._physics_context.enable_gpu_dynamics(True) + self._physics_context.set_broadphase_type("GPU") # Set GPU Pairs capacity and other GPU settings self._physics_context.set_gpu_found_lost_pairs_capacity(gm.GPU_PAIRS_CAPACITY) @@ -485,7 +479,7 @@ def _set_renderer_settings(self): lazy.carb.settings.get_settings().set_bool("/app/renderer/skipMaterialLoading", False) # Below settings are for improving performance: we use the USD / Fabric only for poses. - lazy.carb.settings.get_settings().set_bool("/physics/updateToUsd", not gm.ENABLE_FLATCACHE) + lazy.carb.settings.get_settings().set_bool("/physics/updateToUsd", False) lazy.carb.settings.get_settings().set_bool("/physics/updateParticlesToUsd", False) lazy.carb.settings.get_settings().set_bool("/physics/updateVelocitiesToUsd", False) lazy.carb.settings.get_settings().set_bool("/physics/updateForceSensorsToUsd", False) @@ -1007,9 +1001,7 @@ def play(self): # ignore this if the scale is close to uniform. # We also need to suppress the following error when flat cache is used: # [omni.physx.plugin] Transformation change on non-root links is not supported. - channels = ["omni.usd", "omni.physicsschema.plugin"] - if gm.ENABLE_FLATCACHE: - channels.append("omni.physx.plugin") + channels = ["omni.usd", "omni.physicsschema.plugin", "omni.physx.plugin"] with suppress_omni_log(channels=channels): super().play() @@ -1052,9 +1044,7 @@ def stop(self): if not self.is_stopped(): super().stop() - # If we're using flatcache, we also need to reset its API - if gm.ENABLE_FLATCACHE: - FlatcacheAPI.reset() + FlatcacheAPI.reset() # Run all callbacks for callback in self._callbacks_on_stop.values(): @@ -1702,27 +1692,6 @@ def stage_id(self): """ return self._stage_id - @property - def device(self): - """ - Returns: - device (None or str): Device used in simulation backend - """ - return self._device - - @device.setter - def device(self, device): - """ - Sets the device used for sim backend - - Args: - device (None or str): Device to set for the simulation backend - """ - self._device = device - if self._device is not None and "cuda" in self._device: - device_id = self._settings.get_as_int("/physics/cudaDevice") - self._device = f"cuda:{device_id}" - @property def initial_physics_dt(self): """ diff --git a/omnigibson/systems/macro_particle_system.py b/omnigibson/systems/macro_particle_system.py index a5cb8ef69..1396f400b 100644 --- a/omnigibson/systems/macro_particle_system.py +++ b/omnigibson/systems/macro_particle_system.py @@ -540,9 +540,8 @@ def generate_group_particles( bbox_extents_local = [(self.particle_object.aabb_extent * scale).tolist() for scale in scales] - # If we're using flatcache, we need to update the object's pose on the USD manually - if gm.ENABLE_FLATCACHE: - FlatcacheAPI.sync_raw_object_transforms_in_usd(prim=obj) + # Update the object's pose on the USD manually + FlatcacheAPI.sync_raw_object_transforms_in_usd(prim=obj) # Generate particles z_up = th.zeros((3, 1)) diff --git a/omnigibson/systems/micro_particle_system.py b/omnigibson/systems/micro_particle_system.py index a9322638a..7663883f0 100644 --- a/omnigibson/systems/micro_particle_system.py +++ b/omnigibson/systems/micro_particle_system.py @@ -485,12 +485,6 @@ def initialize(self, scene): # Run super first super().initialize(scene) - # Run sanity checks - if not gm.USE_GPU_DYNAMICS or gm.ENABLE_FLATCACHE: - raise ValueError( - f"Failed to initialize {self.name} system. Please set gm.USE_GPU_DYNAMICS=True and gm.ENABLE_FLATCACHE=False." - ) - self.system_prim = self._create_particle_system() # Get material material = self._get_particle_material_template() @@ -1749,7 +1743,7 @@ def clothify_mesh_prim(self, mesh_prim, remesh=True, particle_distance=None): vertex_normals=new_normals, ) # Apply the inverse of the world transform to get the mesh back into its local frame - tm.apply_transform(th.linalg.inv_ex(scaled_world_transform).inverse) + tm.apply_transform(th.linalg.inv_ex(scaled_world_transform).inverse.cpu()) # Update the mesh prim face_vertex_counts = th.tensor([len(face) for face in tm.faces], dtype=int).cpu().numpy() diff --git a/omnigibson/utils/control_utils.py b/omnigibson/utils/control_utils.py index 7585bbbb1..73fb9a038 100644 --- a/omnigibson/utils/control_utils.py +++ b/omnigibson/utils/control_utils.py @@ -119,7 +119,7 @@ def solve( if isinstance(target_quat, th.Tensor) else th.tensor(target_quat, dtype=th.float64) ) - ik_target_pose = lazy.lula.Pose3(lazy.lula.Rotation3(rot), pos) + ik_target_pose = lazy.lula.Pose3(lazy.lula.Rotation3(rot.cpu()), pos.cpu()) # Set the cspace seed and tolerance initial_joint_pos = self.reset_joint_pos if initial_joint_pos is None else initial_joint_pos diff --git a/omnigibson/utils/processing_utils.py b/omnigibson/utils/processing_utils.py index bb32f35f5..ee8dcb014 100644 --- a/omnigibson/utils/processing_utils.py +++ b/omnigibson/utils/processing_utils.py @@ -1,5 +1,6 @@ import torch as th +import omnigibson as og from omnigibson.utils.python_utils import Serializable @@ -59,7 +60,7 @@ def __init__(self, obs_dim, filter_width): self.obs_dim = obs_dim assert filter_width > 0, f"MovingAverageFilter must have a non-zero size! Got: {filter_width}" self.filter_width = filter_width - self.past_samples = th.zeros((filter_width, obs_dim)) + self.past_samples = th.zeros((filter_width, obs_dim), device=og.sim.device) self.current_idx = 0 self.fully_filled = False # Whether the entire filter buffer is filled or not @@ -103,7 +104,7 @@ def _dump_state(self): state = super()._dump_state() # Add info from this filter - state["past_samples"] = self.past_samples + state["past_samples"] = self.past_samples.cpu() state["current_idx"] = self.current_idx state["fully_filled"] = self.fully_filled @@ -114,7 +115,7 @@ def _load_state(self, state): super()._load_state(state=state) # Load relevant info for this filter - self.past_samples = state["past_samples"] + self.past_samples = state["past_samples"].to(og.sim.device) self.current_idx = state["current_idx"] self.fully_filled = state["fully_filled"] diff --git a/omnigibson/utils/python_utils.py b/omnigibson/utils/python_utils.py index ad6ed3d7f..6d67c29a8 100644 --- a/omnigibson/utils/python_utils.py +++ b/omnigibson/utils/python_utils.py @@ -728,7 +728,10 @@ def nums2array(nums, dim, dtype=th.float32): # Make sure the inputted nums isn't a string assert not isinstance(nums, str), "Only numeric types are supported for this operation!" - out = th.tensor(nums, dtype=dtype) if isinstance(nums, Iterable) else th.ones(dim, dtype=dtype) * nums + if isinstance(nums, th.Tensor): + out = nums.to(dtype=dtype) + else: + out = th.tensor(nums, dtype=dtype) if isinstance(nums, Iterable) else th.ones(dim, dtype=dtype) * nums return out @@ -797,16 +800,20 @@ def h5py_group_to_torch(group): def multi_dim_linspace(start: th.Tensor, stop: th.Tensor, num: int) -> th.Tensor: """ Generate a tensor with evenly spaced values along multiple dimensions. + This function creates a tensor where each slice along the first dimension contains values linearly interpolated between the corresponding elements of 'start' and 'stop'. It's similar to numpy.linspace but works with multi-dimensional inputs in PyTorch. + Args: start (th.Tensor): Starting values for each dimension. stop (th.Tensor): Ending values for each dimension. num (int): Number of samples to generate along the interpolated dimension. + Returns: th.Tensor: A tensor of shape (num, *start.shape) containing the interpolated values. + Example: >>> start = th.tensor([0, 10, 100]) >>> stop = th.tensor([1, 20, 200]) diff --git a/omnigibson/utils/transform_utils.py b/omnigibson/utils/transform_utils.py index fadbc0f6a..7077a0ac7 100644 --- a/omnigibson/utils/transform_utils.py +++ b/omnigibson/utils/transform_utils.py @@ -617,7 +617,7 @@ def pose2mat(pose: Tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor: pos = pos.to(dtype=torch.float32).reshape(3) orn = orn.to(dtype=torch.float32).reshape(4) - homo_pose_mat = torch.eye(4, dtype=torch.float32) + homo_pose_mat = torch.eye(4, dtype=torch.float32, device=pos.device) homo_pose_mat[:3, :3] = quat2mat(orn) homo_pose_mat[:3, 3] = pos @@ -659,7 +659,7 @@ def quat2axisangle(quat): @torch.compile -def axisangle2quat(vec, eps=1e-6): +def axisangle2quat(vec: torch.Tensor, eps: float = 1e-6) -> torch.Tensor: """ Converts scaled axis-angle to quat. Args: @@ -675,17 +675,19 @@ def axisangle2quat(vec, eps=1e-6): vec = vec.reshape(-1, 3) # Grab angle - angle = torch.norm(vec, dim=-1, keepdim=True) + angle = torch.norm(vec, dim=-1, keepdim=True, dtype=torch.float32) # Create return array - quat = torch.zeros(torch.prod(torch.tensor(input_shape, dtype=torch.int)), 4, device=vec.device) + quat = torch.zeros( + torch.prod(torch.tensor(input_shape, dtype=torch.int)), 4, device=vec.device, dtype=torch.float32 + ) quat[:, 3] = 1.0 # Grab indexes where angle is not zero an convert the input to its quaternion form idx = angle.reshape(-1) > eps # torch.nonzero(angle).reshape(-1) quat[idx, :] = torch.cat( [vec[idx, :] * torch.sin(angle[idx, :] / 2.0) / angle[idx, :], torch.cos(angle[idx, :] / 2.0)], dim=-1 - ) + ).float() # Reshape and return output quat = quat.reshape( @@ -742,7 +744,7 @@ def pose_inv(pose_mat): # -t in the original frame, which is -R-1*t in the new frame, and then rotate back by # R-1 to align the axis again. - pose_inv = torch.zeros((4, 4)) + pose_inv = torch.zeros((4, 4), dtype=pose_mat.dtype, device=pose_mat.device) pose_inv[:3, :3] = pose_mat[:3, :3].T pose_inv[:3, 3] = -pose_inv[:3, :3] @ pose_mat[:3, 3] pose_inv[3, 3] = 1.0 diff --git a/omnigibson/utils/ui_utils.py b/omnigibson/utils/ui_utils.py index 4be556d63..b8c9bd994 100644 --- a/omnigibson/utils/ui_utils.py +++ b/omnigibson/utils/ui_utils.py @@ -584,14 +584,15 @@ def __init__(self, robot): self.joint_idx_to_controller = dict() idx = 0 for name, controller in robot._controllers.items(): + dof_idx = controller.dof_idx.tolist() self.controller_info[name] = { "name": type(controller).__name__, "start_idx": idx, - "dofs": controller.dof_idx, + "dofs": dof_idx, "command_dim": controller.command_dim, } idx += controller.command_dim - for i in controller.dof_idx.tolist(): + for i in dof_idx: self.joint_idx_to_controller[i] = controller # Other persistent variables we need to keep track of @@ -728,7 +729,7 @@ def populate_keypress_mapping(self): for i in range(info["command_dim"]): cmd_idx = info["start_idx"] + i self.joint_command_idx.append(cmd_idx) - self.joint_control_idx += info["dofs"].tolist() + self.joint_control_idx += info["dofs"] elif info["name"] == "DifferentialDriveController": self.keypress_mapping[lazy.carb.input.KeyboardInput.I] = {"idx": info["start_idx"] + 0, "val": 0.4} self.keypress_mapping[lazy.carb.input.KeyboardInput.K] = {"idx": info["start_idx"] + 0, "val": -0.4} @@ -745,7 +746,7 @@ def populate_keypress_mapping(self): for i in range(info["command_dim"]): cmd_idx = info["start_idx"] + i self.joint_command_idx.append(cmd_idx) - self.joint_control_idx += info["dofs"].tolist() + self.joint_control_idx += info["dofs"] else: self.keypress_mapping[lazy.carb.input.KeyboardInput.T] = {"idx": info["start_idx"], "val": 1.0} self.gripper_direction[component] = 1.0 diff --git a/omnigibson/utils/usd_utils.py b/omnigibson/utils/usd_utils.py index f7db31d82..0c0ba56e6 100644 --- a/omnigibson/utils/usd_utils.py +++ b/omnigibson/utils/usd_utils.py @@ -633,9 +633,6 @@ def sync_raw_object_transforms_in_usd(cls, prim): prim (EntityPrim): prim whose owned links and joints should have their raw local states updated to match the "true" values found from the dynamic control interface """ - # Make sure flatcache is enabled -- this should NEVER be called otherwise!! - assert gm.ENABLE_FLATCACHE, "Syncing raw object transforms should only occur if flatcache is being used!" - # We're somewhat abusing low-level dynamic control - physx - usd integration, but we (supposedly) know # what we're doing so we suppress logging so we don't see any error messages :D with suppress_omni_log(["omni.physx.plugin"]): @@ -679,9 +676,6 @@ def reset_raw_object_transforms_in_usd(cls, prim): Args: prim (EntityPrim): prim whose owned links and joints should have their local values reset to be zero """ - # Make sure flatcache is enabled -- this should NEVER be called otherwise!! - assert gm.ENABLE_FLATCACHE, "Resetting raw object transforms should only occur if flatcache is being used!" - # We're somewhat abusing low-level dynamic control - physx - usd integration, but we (supposedly) know # what we're doing so we suppress logging so we don't see any error messages :D with suppress_omni_log(["omni.physx.plugin"]): @@ -851,17 +845,17 @@ def clear(self): def flush_control(self): if "dof_position_targets" in self._write_idx_cache: - pos_indices = th.tensor(sorted(self._write_idx_cache["dof_position_targets"])) + pos_indices = th.tensor(sorted(self._write_idx_cache["dof_position_targets"]), device=og.sim.device) pos_targets = self._read_cache["dof_position_targets"] self._view.set_dof_position_targets(pos_targets, pos_indices) if "dof_velocity_targets" in self._write_idx_cache: - vel_indices = th.tensor(sorted(self._write_idx_cache["dof_velocity_targets"])) + vel_indices = th.tensor(sorted(self._write_idx_cache["dof_velocity_targets"]), device=og.sim.device) vel_targets = self._read_cache["dof_velocity_targets"] self._view.set_dof_velocity_targets(vel_targets, vel_indices) if "dof_actuation_forces" in self._write_idx_cache: - eff_indices = th.tensor(sorted(self._write_idx_cache["dof_actuation_forces"])) + eff_indices = th.tensor(sorted(self._write_idx_cache["dof_actuation_forces"]), device=og.sim.device) eff_targets = self._read_cache["dof_actuation_forces"] self._view.set_dof_actuation_forces(eff_targets, eff_indices) @@ -1116,7 +1110,7 @@ def get_jacobian(self, prim_path): def get_relative_jacobian(self, prim_path): jacobian = self.get_jacobian(prim_path) ori_t = T.quat2mat(self.get_position_orientation(prim_path)[1]).T - tf = th.zeros((1, 6, 6), dtype=th.float32) + tf = th.zeros((1, 6, 6), dtype=th.float32, device=og.sim.device) tf[:, :3, :3] = ori_t tf[:, 3:, 3:] = ori_t return tf @ jacobian @@ -1384,8 +1378,8 @@ def mesh_prim_mesh_to_trimesh_mesh(mesh_prim, include_normals=True, include_texc mesh_type = mesh_prim.GetPrimTypeInfo().GetTypeName() assert mesh_type == "Mesh", f"Expected mesh prim to have type Mesh, got {mesh_type}" face_vertex_counts = vtarray_to_torch(mesh_prim.GetAttribute("faceVertexCounts").Get(), dtype=th.int) - vertices = vtarray_to_torch(mesh_prim.GetAttribute("points").Get()) - face_indices = vtarray_to_torch(mesh_prim.GetAttribute("faceVertexIndices").Get(), dtype=th.int) + vertices = vtarray_to_torch(mesh_prim.GetAttribute("points").Get()).cpu() + face_indices = vtarray_to_torch(mesh_prim.GetAttribute("faceVertexIndices").Get(), dtype=th.int).cpu() faces = [] i = 0 @@ -1397,7 +1391,7 @@ def mesh_prim_mesh_to_trimesh_mesh(mesh_prim, include_normals=True, include_texc kwargs = dict(vertices=vertices, faces=faces) if include_normals: - kwargs["vertex_normals"] = vtarray_to_torch(mesh_prim.GetAttribute("normals").Get()) + kwargs["vertex_normals"] = vtarray_to_torch(mesh_prim.GetAttribute("normals").Get()).cpu() if include_texcoord: raw_texture = mesh_prim.GetAttribute("primvars:st").Get() @@ -1462,7 +1456,7 @@ def mesh_prim_to_trimesh_mesh(mesh_prim, include_normals=True, include_texcoord= trimesh_mesh = mesh_prim_shape_to_trimesh_mesh(mesh_prim) if world_frame: - trimesh_mesh.apply_transform(PoseAPI.get_world_pose_with_scale(mesh_prim.GetPath().pathString)) + trimesh_mesh.apply_transform(PoseAPI.get_world_pose_with_scale(mesh_prim.GetPath().pathString).cpu()) return trimesh_mesh diff --git a/tests/benchmark/benchmark_interactive_scene.py b/tests/benchmark/benchmark_interactive_scene.py index 5f887f10b..110b5ff04 100644 --- a/tests/benchmark/benchmark_interactive_scene.py +++ b/tests/benchmark/benchmark_interactive_scene.py @@ -22,8 +22,6 @@ gm.HEADLESS = False gm.GUI_VIEWPORT_ONLY = True gm.RENDER_VIEWER_CAMERA = False -gm.ENABLE_FLATCACHE = True -gm.USE_GPU_DYNAMICS = False gm.ENABLE_OBJECT_STATES = False gm.ENABLE_TRANSITION_RULES = False gm.DEFAULT_VIEWER_WIDTH = 128 diff --git a/tests/benchmark/profiling.py b/tests/benchmark/profiling.py index eb449d9f2..a18384956 100644 --- a/tests/benchmark/profiling.py +++ b/tests/benchmark/profiling.py @@ -19,7 +19,6 @@ parser.add_argument("-s", "--scene", default="") parser.add_argument("-c", "--cloth", action="store_true") parser.add_argument("-w", "--fluids", action="store_true") -parser.add_argument("-g", "--gpu_denamics", action="store_true") parser.add_argument("-p", "--macro_particle_system", action="store_true") PROFILING_FIELDS = ["FPS", "Omni step time", "Non-omni step time", "Memory usage", "Vram usage"] @@ -42,8 +41,6 @@ def main(): gm.ENABLE_HQ_RENDERING = args.fluids gm.ENABLE_OBJECT_STATES = True gm.ENABLE_TRANSITION_RULES = True - gm.ENABLE_FLATCACHE = not args.cloth - gm.USE_GPU_DYNAMICS = args.gpu_denamics cfg = { "env": { diff --git a/tests/conftest.py b/tests/conftest.py index c9162fe65..eb92f499b 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,5 +1,12 @@ +import pytest + import omnigibson as og def pytest_unconfigure(config): og.shutdown() + + +@pytest.fixture(params=["cpu", "cuda"]) +def pipeline_mode(request): + return request.param diff --git a/tests/test_controllers.py b/tests/test_controllers.py index 7b13674aa..f9411a927 100644 --- a/tests/test_controllers.py +++ b/tests/test_controllers.py @@ -7,9 +7,13 @@ from omnigibson.robots import LocomotionRobot -def test_arm_control(): +@pytest.mark.parametrize("pipeline_mode", ["cpu", "cuda"], indirect=True) +def test_arm_control(pipeline_mode): # Create env cfg = { + "env": { + "device": pipeline_mode, + }, "scene": { "type": "Scene", }, @@ -296,12 +300,12 @@ def check_ori_error(curr_orientation, init_orientation, tol=0.1): curr_pos, curr_quat = robot.get_relative_eef_pose(arm=arm) arm_controller = robot.controllers[f"arm_{arm}"] arm_goal = arm_controller.goal - target_pos = arm_goal["target_pos"] + target_pos = arm_goal["target_pos"].cpu() target_quat = ( arm_goal["target_quat"] if controller == "InverseKinematicsController" else T.mat2quat(arm_goal["target_ori_mat"]) - ) + ).cpu() pos_check = err_checks[controller_mode][action_name]["pos"] if pos_check is not None: is_valid_pos = pos_check(target_pos, curr_pos, init_pos) diff --git a/tests/test_data_collection.py b/tests/test_data_collection.py index 5f8d0b3df..e2573437c 100644 --- a/tests/test_data_collection.py +++ b/tests/test_data_collection.py @@ -9,10 +9,12 @@ from omnigibson.objects import DatasetObject -def test_data_collect_and_playback(): +@pytest.mark.parametrize("pipeline_mode", ["cpu", "cuda"], indirect=True) +def test_data_collect_and_playback(pipeline_mode): cfg = { "env": { "external_sensors": [], + "device": pipeline_mode, }, "scene": { "type": "InteractiveTraversableScene", @@ -37,7 +39,6 @@ def test_data_collect_and_playback(): if og.sim is None: # Make sure GPU dynamics are enabled (GPU dynamics needed for cloth) gm.ENABLE_OBJECT_STATES = True - gm.USE_GPU_DYNAMICS = True gm.ENABLE_TRANSITION_RULES = False else: # Make sure sim is stopped @@ -59,7 +60,7 @@ def test_data_collect_and_playback(): for i in range(2): env.reset() for _ in range(5): - env.step(env.robots[0].action_space.sample()) + env.step(th.tensor(env.robots[0].action_space.sample())) # Manually add a random object, e.g.: a banana, and place on the floor obj = DatasetObject(name="banana", category="banana") @@ -68,14 +69,14 @@ def test_data_collect_and_playback(): # Take a few more steps for _ in range(5): - env.step(env.robots[0].action_space.sample()) + env.step(th.tensor(env.robots[0].action_space.sample())) # Manually remove the added object env.scene.remove_object(obj) # Take a few more steps for _ in range(5): - env.step(env.robots[0].action_space.sample()) + env.step(th.tensor(env.robots[0].action_space.sample())) # Add water particles water = env.scene.get_system("water") @@ -84,14 +85,14 @@ def test_data_collect_and_playback(): # Take a few more steps for _ in range(5): - env.step(env.robots[0].action_space.sample()) + env.step(th.tensor(env.robots[0].action_space.sample())) # Clear the system env.scene.clear_system("water") # Take a few more steps for _ in range(5): - env.step(env.robots[0].action_space.sample()) + env.step(th.tensor(env.robots[0].action_space.sample())) # Save this data env.save_data() diff --git a/tests/test_dump_load_states.py b/tests/test_dump_load_states.py index c31f6c7eb..a29a0ab3f 100644 --- a/tests/test_dump_load_states.py +++ b/tests/test_dump_load_states.py @@ -7,7 +7,7 @@ @og_test -def test_dump_load(env): +def test_dump_load(env, pipeline_mode): breakfast_table = env.scene.object_registry("name", "breakfast_table") for system_name, system_class in SYSTEM_EXAMPLES.items(): system = env.scene.get_system(system_name) @@ -27,7 +27,7 @@ def test_dump_load(env): @og_test -def test_dump_load_serialized(env): +def test_dump_load_serialized(env, pipeline_mode): breakfast_table = env.scene.object_registry("name", "breakfast_table") for system_name, system_class in SYSTEM_EXAMPLES.items(): system = env.scene.get_system(system_name) @@ -46,7 +46,7 @@ def test_dump_load_serialized(env): @og_test -def test_save_restore_partial(env): +def test_save_restore_partial(env, pipeline_mode): breakfast_table = env.scene.object_registry("name", "breakfast_table") decrypted_fd, tmp_json_path = tempfile.mkstemp("test_save_restore.json", dir=og.tempdir) @@ -65,7 +65,7 @@ def test_save_restore_partial(env): @og_test -def test_save_restore_full(env): +def test_save_restore_full(env, pipeline_mode): decrypted_fd, tmp_json_path = tempfile.mkstemp("test_save_restore.json", dir=og.tempdir) og.sim.save([tmp_json_path]) diff --git a/tests/test_envs.py b/tests/test_envs.py index 33b56a1a8..f789052ef 100644 --- a/tests/test_envs.py +++ b/tests/test_envs.py @@ -1,69 +1,74 @@ import pytest +import torch as th import omnigibson as og from omnigibson.macros import gm -def task_tester(task_type): - cfg = { - "scene": { - "type": "InteractiveTraversableScene", - "scene_model": "Rs_int", - "load_object_categories": ["floors", "breakfast_table"], - }, - "robots": [ - { - "type": "Fetch", - "obs_modalities": [], - } - ], - # Task kwargs - "task": { - "type": task_type, - # BehaviorTask-specific - "activity_name": "laying_wood_floors", - "online_object_sampling": True, - }, - } +@pytest.mark.parametrize("pipeline_mode", ["cpu", "cuda"], indirect=True) +class TestTasks: + def task_tester(self, task_type, pipeline_mode): + cfg = { + "env": { + "device": pipeline_mode, + }, + "scene": { + "type": "InteractiveTraversableScene", + "scene_model": "Rs_int", + "load_object_categories": ["floors", "breakfast_table"], + }, + "robots": [ + { + "type": "Fetch", + "obs_modalities": [], + } + ], + # Task kwargs + "task": { + "type": task_type, + # BehaviorTask-specific + "activity_name": "laying_wood_floors", + "online_object_sampling": True, + }, + } if og.sim is None: # Make sure GPU dynamics are enabled (GPU dynamics needed for cloth) gm.ENABLE_OBJECT_STATES = True - gm.USE_GPU_DYNAMICS = True gm.ENABLE_TRANSITION_RULES = False else: # Make sure sim is stopped og.sim.stop() - # Create the environment - env = og.Environment(configs=cfg) + # Create the environment + env = og.Environment(configs=cfg) - env.reset() - for _ in range(5): - env.step(env.robots[0].action_space.sample()) + env.reset() + for _ in range(5): + env.step(th.tensor(env.robots[0].action_space.sample(), dtype=th.float32)) - # Clear the sim - og.clear() - - -def test_dummy_task(): - task_tester("DummyTask") + # Clear the sim + og.clear() + def test_dummy_task(self, pipeline_mode): + self.task_tester("DummyTask", pipeline_mode) -def test_point_reaching_task(): - task_tester("PointReachingTask") + def test_point_reaching_task(self, pipeline_mode): + self.task_tester("PointReachingTask", pipeline_mode) + def test_point_navigation_task(self, pipeline_mode): + self.task_tester("PointNavigationTask", pipeline_mode) -def test_point_navigation_task(): - task_tester("PointNavigationTask") + def test_behavior_task(self, pipeline_mode): + self.task_tester("BehaviorTask", pipeline_mode) -def test_behavior_task(): - task_tester("BehaviorTask") - - -def test_rs_int_full_load(): +@pytest.mark.parametrize("pipeline_mode", ["cpu", "cuda"], indirect=True) +def test_rs_int_full_load(pipeline_mode): cfg = { + "env": { + "device": pipeline_mode, + }, "scene": { "type": "InteractiveTraversableScene", "scene_model": "Rs_int", @@ -89,7 +94,7 @@ def test_rs_int_full_load(): env.reset() for _ in range(5): - env.step(env.robots[0].action_space.sample()) + env.step(th.tensor(env.robots[0].action_space.sample(), dtype=th.float32)) # Clear the sim og.clear() diff --git a/tests/test_multiple_envs.py b/tests/test_multiple_envs.py index 5ee8b6d1b..888c8fc58 100644 --- a/tests/test_multiple_envs.py +++ b/tests/test_multiple_envs.py @@ -10,480 +10,473 @@ from omnigibson.utils.transform_utils import quat_multiply -def setup_multi_environment(num_of_envs, robot="Fetch", additional_objects_cfg=[]): - cfg = { - "scene": { - "type": "InteractiveTraversableScene", - "scene_model": "Rs_int", - "load_object_categories": ["floors", "walls"], - }, - "robots": [ - { - "type": robot, - "obs_modalities": [], - } - ], - } - - cfg["objects"] = additional_objects_cfg - - if og.sim is None: - # Make sure GPU dynamics are enabled (GPU dynamics needed for cloth) - gm.RENDER_VIEWER_CAMERA = False - gm.ENABLE_OBJECT_STATES = True - gm.USE_GPU_DYNAMICS = True - gm.ENABLE_FLATCACHE = False - gm.ENABLE_TRANSITION_RULES = False - else: - # Make sure sim is stopped - og.sim.stop() - - vec_env = og.VectorEnvironment(num_of_envs, cfg) - return vec_env - - -def test_multi_scene_dump_load_states(): - vec_env = setup_multi_environment(3) - robot_0 = vec_env.envs[0].scene.robots[0] - robot_1 = vec_env.envs[1].scene.robots[0] - robot_2 = vec_env.envs[2].scene.robots[0] - - robot_0_pos = robot_0.get_position_orientation()[0] - robot_1_pos = robot_1.get_position_orientation()[0] - robot_2_pos = robot_2.get_position_orientation()[0] - - dist_0_1 = robot_1_pos - robot_0_pos - dist_1_2 = robot_2_pos - robot_1_pos - - assert th.allclose(dist_0_1, dist_1_2, atol=1e-3) - - # Set different poses for the cube in each environment - pose_1 = (th.tensor([1, 1, 1], dtype=th.float32), th.tensor([0, 0, 0, 1], dtype=th.float32)) - pose_2 = (th.tensor([0, 2, 1], dtype=th.float32), th.tensor([0, 0, 0.7071, 0.7071], dtype=th.float32)) - pose_3 = (th.tensor([-1, -1, 0.5], dtype=th.float32), th.tensor([0.5, 0.5, 0.5, 0.5], dtype=th.float32)) - - robot_0.set_position_orientation(*pose_1, frame="scene") - robot_1.set_position_orientation(*pose_2, frame="scene") - robot_2.set_position_orientation(*pose_3, frame="scene") - - # Run simulation for a bit - for _ in range(10): - og.sim.step() - - initial_robot_pos_scene_1 = robot_1.get_position_orientation(frame="scene") - initial_robot_pos_scene_2 = robot_2.get_position_orientation(frame="scene") - initial_robot_pos_scene_0 = robot_0.get_position_orientation(frame="scene") - - # Save states - robot_0_state = vec_env.envs[0].scene._dump_state() - robot_1_state = vec_env.envs[1].scene._dump_state() - robot_2_state = vec_env.envs[2].scene._dump_state() - - # reset the object_positions - vec_env.reset() - - # Load the states in a different order - vec_env.envs[1].scene._load_state(robot_1_state) - vec_env.envs[2].scene._load_state(robot_2_state) - vec_env.envs[0].scene._load_state(robot_0_state) - - post_robot_pos_scene_1 = vec_env.envs[1].scene.robots[0].get_position_orientation(frame="scene") - post_robot_pos_scene_2 = vec_env.envs[2].scene.robots[0].get_position_orientation(frame="scene") - post_robot_pos_scene_0 = vec_env.envs[0].scene.robots[0].get_position_orientation(frame="scene") - - # Check that the poses are the same - assert th.allclose(initial_robot_pos_scene_0[0], post_robot_pos_scene_0[0], atol=1e-3) - assert th.allclose(initial_robot_pos_scene_1[0], post_robot_pos_scene_1[0], atol=1e-3) - assert th.allclose(initial_robot_pos_scene_2[0], post_robot_pos_scene_2[0], atol=1e-3) - - assert th.allclose(initial_robot_pos_scene_0[1], post_robot_pos_scene_0[1], atol=1e-3) - assert th.allclose(initial_robot_pos_scene_1[1], post_robot_pos_scene_1[1], atol=1e-3) - assert th.allclose(initial_robot_pos_scene_2[1], post_robot_pos_scene_2[1], atol=1e-3) - - og.clear() - - -def test_multi_scene_get_local_position(): - vec_env = setup_multi_environment(3) - - robot_1_pos_local = vec_env.envs[1].scene.robots[0].get_position_orientation(frame="scene")[0] - robot_1_pos_global = vec_env.envs[1].scene.robots[0].get_position_orientation()[0] - - pos_scene = vec_env.envs[1].scene.get_position_orientation()[0] - - assert th.allclose(robot_1_pos_global, pos_scene + robot_1_pos_local, atol=1e-3) - og.clear() - - -def test_multi_scene_set_local_position(): - vec_env = setup_multi_environment(3) - - # Get the robot from the second environment - robot = vec_env.envs[1].scene.robots[0] - - # Get the initial global position of the robot - initial_global_pos = robot.get_position_orientation()[0] - - # Define a new global position - new_global_pos = initial_global_pos + th.tensor([1.0, 0.5, 0.0], dtype=th.float32) - - # Set the new global position - robot.set_position_orientation(position=new_global_pos) - - # Get the updated global position - updated_global_pos = robot.get_position_orientation()[0] - - # Get the scene's global position - scene_pos = vec_env.envs[1].scene.get_position_orientation()[0] - - # Get the updated local position - updated_local_pos = robot.get_position_orientation(frame="scene")[0] +@pytest.mark.parametrize("pipeline_mode", ["cpu", "cuda"], indirect=True) +class TestMultiEnvironment: + + def setup_multi_environment(self, pipeline_mode, num_of_envs, robot="Fetch", additional_objects_cfg=[]): + cfg = { + "env": {"device": pipeline_mode}, + "scene": { + "type": "InteractiveTraversableScene", + "scene_model": "Rs_int", + "load_object_categories": ["floors", "walls"], + }, + "robots": [ + { + "type": robot, + "obs_modalities": [], + } + ], + } - # Calculate expected local position - expected_local_pos = new_global_pos - scene_pos + cfg["objects"] = additional_objects_cfg - # Assert that the global position has been updated correctly - assert th.allclose( - updated_global_pos, new_global_pos, atol=1e-3 - ), f"Updated global position {updated_global_pos} does not match expected {new_global_pos}" + if og.sim is None: + # Make sure GPU dynamics are enabled (GPU dynamics needed for cloth) + gm.RENDER_VIEWER_CAMERA = False + gm.ENABLE_OBJECT_STATES = True + gm.ENABLE_TRANSITION_RULES = False + else: + # Make sure sim is stopped + og.sim.stop() - # Assert that the local position has been updated correctly - assert th.allclose( - updated_local_pos, expected_local_pos, atol=1e-3 - ), f"Updated local position {updated_local_pos} does not match expected {expected_local_pos}" + vec_env = og.VectorEnvironment(num_of_envs, cfg) + return vec_env - # Assert that the change in global position is correct - global_pos_change = updated_global_pos - initial_global_pos - expected_change = th.tensor([1.0, 0.5, 0.0], dtype=th.float32) - assert th.allclose( - global_pos_change, expected_change, atol=1e-3 - ), f"Global position change {global_pos_change} does not match expected change {expected_change}" + def test_multi_scene_dump_load_states(self, pipeline_mode): + vec_env = self.setup_multi_environment(pipeline_mode, 3) + robot_0 = vec_env.envs[0].scene.robots[0] + robot_1 = vec_env.envs[1].scene.robots[0] + robot_2 = vec_env.envs[2].scene.robots[0] - og.clear() + robot_0_pos = robot_0.get_position_orientation()[0] + robot_1_pos = robot_1.get_position_orientation()[0] + robot_2_pos = robot_2.get_position_orientation()[0] + dist_0_1 = robot_1_pos - robot_0_pos + dist_1_2 = robot_2_pos - robot_1_pos -def test_multi_scene_scene_prim(): - vec_env = setup_multi_environment(1) - original_robot_pos = vec_env.envs[0].scene.robots[0].get_position_orientation()[0] - scene_state = vec_env.envs[0].scene._dump_state() - scene_prim_displacement = th.tensor([10.0, 0.0, 0.0], dtype=th.float32) - original_scene_prim_pos = vec_env.envs[0].scene._scene_prim.get_position_orientation()[0] - vec_env.envs[0].scene.set_position_orientation(position=original_scene_prim_pos + scene_prim_displacement) - vec_env.envs[0].scene._load_state(scene_state) - new_scene_prim_pos = vec_env.envs[0].scene._scene_prim.get_position_orientation()[0] - new_robot_pos = vec_env.envs[0].scene.robots[0].get_position_orientation()[0] - assert th.allclose(new_scene_prim_pos - original_scene_prim_pos, scene_prim_displacement, atol=1e-3) - assert th.allclose(new_robot_pos - original_robot_pos, scene_prim_displacement, atol=1e-3) + assert th.allclose(dist_0_1, dist_1_2, atol=1e-3) - og.clear() + # Set different poses for the cube in each environment + pose_1 = (th.tensor([1, 1, 1], dtype=th.float32), th.tensor([0, 0, 0, 1], dtype=th.float32)) + pose_2 = (th.tensor([0, 2, 1], dtype=th.float32), th.tensor([0, 0, 0.7071, 0.7071], dtype=th.float32)) + pose_3 = (th.tensor([-1, -1, 0.5], dtype=th.float32), th.tensor([0.5, 0.5, 0.5, 0.5], dtype=th.float32)) + robot_0.set_position_orientation(*pose_1, frame="scene") + robot_1.set_position_orientation(*pose_2, frame="scene") + robot_2.set_position_orientation(*pose_3, frame="scene") -def test_multi_scene_particle_source(): - sink_cfg = dict( - type="DatasetObject", - name="sink", - category="sink", - model="egwapq", - bounding_box=[2.427, 0.625, 1.2], - abilities={ - "toggleable": {}, - "particleSource": { - "conditions": { - "water": [ - (ParticleModifyCondition.TOGGLEDON, True) - ], # Must be toggled on for water source to be active + # Run simulation for a bit + for _ in range(10): + og.sim.step() + + initial_robot_pos_scene_1 = robot_1.get_position_orientation(frame="scene") + initial_robot_pos_scene_2 = robot_2.get_position_orientation(frame="scene") + initial_robot_pos_scene_0 = robot_0.get_position_orientation(frame="scene") + + # Save states + robot_0_state = vec_env.envs[0].scene._dump_state() + robot_1_state = vec_env.envs[1].scene._dump_state() + robot_2_state = vec_env.envs[2].scene._dump_state() + og.clear() + + # recreate the environments + vec_env = self.setup_multi_environment(pipeline_mode, 3) + + # Load the states in a different order + vec_env.envs[1].scene._load_state(robot_1_state) + vec_env.envs[2].scene._load_state(robot_2_state) + vec_env.envs[0].scene._load_state(robot_0_state) + + post_robot_pos_scene_1 = vec_env.envs[1].scene.robots[0].get_position_orientation(frame="scene") + post_robot_pos_scene_2 = vec_env.envs[2].scene.robots[0].get_position_orientation(frame="scene") + post_robot_pos_scene_0 = vec_env.envs[0].scene.robots[0].get_position_orientation(frame="scene") + + # Check that the poses are the same + assert th.allclose(initial_robot_pos_scene_0[0], post_robot_pos_scene_0[0], atol=1e-3) + assert th.allclose(initial_robot_pos_scene_1[0], post_robot_pos_scene_1[0], atol=1e-3) + assert th.allclose(initial_robot_pos_scene_2[0], post_robot_pos_scene_2[0], atol=1e-3) + + assert th.allclose(initial_robot_pos_scene_0[1], post_robot_pos_scene_0[1], atol=1e-3) + assert th.allclose(initial_robot_pos_scene_1[1], post_robot_pos_scene_1[1], atol=1e-3) + assert th.allclose(initial_robot_pos_scene_2[1], post_robot_pos_scene_2[1], atol=1e-3) + + og.clear() + + def test_multi_scene_get_local_position(self, pipeline_mode): + vec_env = self.setup_multi_environment(pipeline_mode, 3) + + robot_1_pos_local = vec_env.envs[1].scene.robots[0].get_position_orientation(frame="scene")[0] + robot_1_pos_global = vec_env.envs[1].scene.robots[0].get_position_orientation()[0] + + pos_scene = vec_env.envs[1].scene.get_position_orientation()[0] + + assert th.allclose(robot_1_pos_global, pos_scene + robot_1_pos_local, atol=1e-3) + og.clear() + + def test_multi_scene_set_local_position(self, pipeline_mode): + vec_env = self.setup_multi_environment(pipeline_mode, 3) + + # Get the robot from the second environment + robot = vec_env.envs[1].scene.robots[0] + + # Get the initial global position of the robot + initial_global_pos = robot.get_position_orientation()[0] + + # Define a new global position + new_global_pos = initial_global_pos + th.tensor([1.0, 0.5, 0.0], dtype=th.float32) + + # Set the new global position + robot.set_position_orientation(position=new_global_pos) + + # Get the updated global position + updated_global_pos = robot.get_position_orientation()[0] + + # Get the scene's global position + scene_pos = vec_env.envs[1].scene.get_position_orientation()[0] + + # Get the updated local position + updated_local_pos = robot.get_position_orientation(frame="scene")[0] + + # Calculate expected local position + expected_local_pos = new_global_pos - scene_pos + + # Assert that the global position has been updated correctly + assert th.allclose( + updated_global_pos, new_global_pos, atol=1e-3 + ), f"Updated global position {updated_global_pos} does not match expected {new_global_pos}" + + # Assert that the local position has been updated correctly + assert th.allclose( + updated_local_pos, expected_local_pos, atol=1e-3 + ), f"Updated local position {updated_local_pos} does not match expected {expected_local_pos}" + + # Assert that the change in global position is correct + global_pos_change = updated_global_pos - initial_global_pos + expected_change = th.tensor([1.0, 0.5, 0.0], dtype=th.float32) + assert th.allclose( + global_pos_change, expected_change, atol=1e-3 + ), f"Global position change {global_pos_change} does not match expected change {expected_change}" + + og.clear() + + def test_multi_scene_scene_prim(self, pipeline_mode): + vec_env = self.setup_multi_environment(pipeline_mode, 1) + original_robot_pos = vec_env.envs[0].scene.robots[0].get_position_orientation()[0] + scene_state = vec_env.envs[0].scene._dump_state() + scene_prim_displacement = th.tensor([10.0, 0.0, 0.0], dtype=th.float32) + original_scene_prim_pos = vec_env.envs[0].scene._scene_prim.get_position_orientation()[0] + vec_env.envs[0].scene.set_position_orientation(position=original_scene_prim_pos + scene_prim_displacement) + vec_env.envs[0].scene._load_state(scene_state) + new_scene_prim_pos = vec_env.envs[0].scene._scene_prim.get_position_orientation()[0] + new_robot_pos = vec_env.envs[0].scene.robots[0].get_position_orientation()[0] + assert th.allclose(new_scene_prim_pos - original_scene_prim_pos, scene_prim_displacement, atol=1e-3) + assert th.allclose(new_robot_pos - original_robot_pos, scene_prim_displacement, atol=1e-3) + + og.clear() + + def test_multi_scene_particle_source(self, pipeline_mode): + sink_cfg = dict( + type="DatasetObject", + name="sink", + category="sink", + model="egwapq", + bounding_box=[2.427, 0.625, 1.2], + abilities={ + "toggleable": {}, + "particleSource": { + "conditions": { + "water": [ + (ParticleModifyCondition.TOGGLEDON, True) + ], # Must be toggled on for water source to be active + }, + "initial_speed": 0.0, # Water merely falls out of the spout }, - "initial_speed": 0.0, # Water merely falls out of the spout - }, - "particleSink": { - "conditions": { - "water": [], # No conditions, always sinking nearby particles + "particleSink": { + "conditions": { + "water": [], # No conditions, always sinking nearby particles + }, }, }, - }, - position=[0.0, -1.5, 0.42], - ) - - vec_env = setup_multi_environment(3, additional_objects_cfg=[sink_cfg]) - - for env in vec_env.envs: - sink = env.scene.object_registry("name", "sink") - assert sink.states[object_states.ToggledOn].set_value(True) - - for _ in range(50): - og.sim.step() - - og.clear() - + position=[0.0, -1.5, 0.42], + ) -def test_multi_scene_position_orientation_relative_to_scene(): - vec_env = setup_multi_environment(3) + vec_env = self.setup_multi_environment(pipeline_mode, 3, additional_objects_cfg=[sink_cfg]) - # Get the robot from the second environment - robot = vec_env.envs[1].scene.robots[0] + for env in vec_env.envs: + sink = env.scene.object_registry("name", "sink") + assert sink.states[object_states.ToggledOn].set_value(True) - # Define a new position and orientation relative to the scene - new_relative_pos = th.tensor([1.0, 2.0, 0.5]) - new_relative_ori = th.tensor([0, 0, 0.7071, 0.7071]) # 90 degrees rotation around z-axis + for _ in range(50): + og.sim.step() - # Set the new position and orientation relative to the scene - robot.set_position_orientation(position=new_relative_pos, orientation=new_relative_ori, frame="scene") + og.clear() - # Get the updated position and orientation relative to the scene - updated_relative_pos, updated_relative_ori = robot.get_position_orientation(frame="scene") + def test_multi_scene_position_orientation_relative_to_scene(self, pipeline_mode): + vec_env = self.setup_multi_environment(pipeline_mode, 3) - # Assert that the relative position has been updated correctly - assert th.allclose( - updated_relative_pos, new_relative_pos, atol=1e-3 - ), f"Updated relative position {updated_relative_pos} does not match expected {new_relative_pos}" + # Get the robot from the second environment + robot = vec_env.envs[1].scene.robots[0] - # Assert that the relative orientation has been updated correctly - assert th.allclose( - updated_relative_ori, new_relative_ori, atol=1e-3 - ), f"Updated relative orientation {updated_relative_ori} does not match expected {new_relative_ori}" + # Define a new position and orientation relative to the scene + new_relative_pos = th.tensor([1.0, 2.0, 0.5]) + new_relative_ori = th.tensor([0, 0, 0.7071, 0.7071]) # 90 degrees rotation around z-axis - # Get the scene's global position and orientation - scene_pos, scene_ori = vec_env.envs[1].scene.get_position_orientation() + # Set the new position and orientation relative to the scene + robot.set_position_orientation(position=new_relative_pos, orientation=new_relative_ori, frame="scene") - # Get the robot's global position and orientation - global_pos, global_ori = robot.get_position_orientation() + # Get the updated position and orientation relative to the scene + updated_relative_pos, updated_relative_ori = robot.get_position_orientation(frame="scene") - # Calculate expected global position - expected_global_pos = scene_pos + updated_relative_pos + # Assert that the relative position has been updated correctly + assert th.allclose( + updated_relative_pos, new_relative_pos, atol=1e-3 + ), f"Updated relative position {updated_relative_pos} does not match expected {new_relative_pos}" - # Assert that the global position is correct - assert th.allclose( - global_pos, expected_global_pos, atol=1e-3 - ), f"Global position {global_pos} does not match expected {expected_global_pos}" + # Assert that the relative orientation has been updated correctly + assert th.allclose( + updated_relative_ori, new_relative_ori, atol=1e-3 + ), f"Updated relative orientation {updated_relative_ori} does not match expected {new_relative_ori}" - # Calculate expected global orientation - expected_global_ori = quat_multiply(scene_ori, new_relative_ori) + # Get the scene's global position and orientation + scene_pos, scene_ori = vec_env.envs[1].scene.get_position_orientation() - # Assert that the global orientation is correct - assert th.allclose( - global_ori, expected_global_ori, atol=1e-3 - ), f"Global orientation {global_ori} does not match expected {expected_global_ori}" + # Get the robot's global position and orientation + global_pos, global_ori = robot.get_position_orientation() - og.clear() + # Calculate expected global position + expected_global_pos = scene_pos + updated_relative_pos + # Assert that the global position is correct + assert th.allclose( + global_pos, expected_global_pos, atol=1e-3 + ), f"Global position {global_pos} does not match expected {expected_global_pos}" -def test_tiago_getter(): - vec_env = setup_multi_environment(2, robot="Tiago") - robot1 = vec_env.envs[0].scene.robots[0] + # Calculate expected global orientation + expected_global_ori = quat_multiply(scene_ori, new_relative_ori) - robot1_world_position, robot1_world_orientation = robot1.get_position_orientation() - robot1_scene_position, robot1_scene_orientation = robot1.get_position_orientation(frame="scene") + # Assert that the global orientation is correct + assert th.allclose( + global_ori, expected_global_ori, atol=1e-3 + ), f"Global orientation {global_ori} does not match expected {expected_global_ori}" - # Test the get_position_orientation method for 3 different frames - # since the robot is at the origin, the position and orientation should be the same - assert th.allclose(robot1_world_position, robot1_scene_position, atol=1e-3) - assert th.allclose(robot1_world_orientation, robot1_scene_orientation, atol=1e-3) + og.clear() - # test if the scene position is non-zero, the getter with scene and world frame should return different values - robot2 = vec_env.envs[1].scene.robots[0] - scene_position, scene_orientation = vec_env.envs[1].scene.get_position_orientation() + def test_tiago_getter(self, pipeline_mode): + vec_env = self.setup_multi_environment(pipeline_mode, 2, robot="Tiago") + robot1 = vec_env.envs[0].scene.robots[0] - robot2_world_position, robot2_world_orientation = robot2.get_position_orientation() - robot2_scene_position, robot2_scene_orientation = robot2.get_position_orientation(frame="scene") + robot1_world_position, robot1_world_orientation = robot1.get_position_orientation() + robot1_scene_position, robot1_scene_orientation = robot1.get_position_orientation(frame="scene") - combined_position, combined_orientation = T.pose_transform( - scene_position, scene_orientation, robot2_scene_position, robot2_scene_orientation - ) - assert th.allclose(robot2_world_position, combined_position, atol=1e-3) - assert th.allclose(robot2_world_orientation, combined_orientation, atol=1e-3) + # Test the get_position_orientation method for 3 different frames + # since the robot is at the origin, the position and orientation should be the same + assert th.allclose(robot1_world_position, robot1_scene_position, atol=1e-3) + assert th.allclose(robot1_world_orientation, robot1_scene_orientation, atol=1e-3) - # Clean up - og.clear() + # test if the scene position is non-zero, the getter with scene and world frame should return different values + robot2 = vec_env.envs[1].scene.robots[0] + scene_position, scene_orientation = vec_env.envs[1].scene.get_position_orientation() + robot2_world_position, robot2_world_orientation = robot2.get_position_orientation() + robot2_scene_position, robot2_scene_orientation = robot2.get_position_orientation(frame="scene") -def test_tiago_setter(): - vec_env = setup_multi_environment(2, robot="Tiago") + combined_position, combined_orientation = T.pose_transform( + scene_position, scene_orientation, robot2_scene_position, robot2_scene_orientation + ) + assert th.allclose(robot2_world_position, combined_position, atol=1e-3) + assert th.allclose(robot2_world_orientation, combined_orientation, atol=1e-3) - # use a robot with non-zero scene position - robot = vec_env.envs[1].scene.robots[0] + # Clean up + og.clear() - # Test setting position and orientation in world frame - new_world_pos = th.tensor([1.0, 2.0, 0.5]) - new_world_ori = T.euler2quat(th.tensor([0, 0, th.pi / 2])) - robot.set_position_orientation(position=new_world_pos, orientation=new_world_ori) + def test_tiago_setter(self, pipeline_mode): + vec_env = self.setup_multi_environment(pipeline_mode, 2, robot="Tiago") - got_world_pos, got_world_ori = robot.get_position_orientation() - assert th.allclose(got_world_pos, new_world_pos, atol=1e-3) - assert th.allclose(got_world_ori, new_world_ori, atol=1e-3) + # use a robot with non-zero scene position + robot = vec_env.envs[1].scene.robots[0] - # Test setting position and orientation in scene frame - new_scene_pos = th.tensor([0.5, 1.0, 0.25]) - new_scene_ori = T.euler2quat(th.tensor([0, th.pi / 4, 0])) - robot.set_position_orientation(position=new_scene_pos, orientation=new_scene_ori, frame="scene") + # Test setting position and orientation in world frame + new_world_pos = th.tensor([1.0, 2.0, 0.5]) + new_world_ori = T.euler2quat(th.tensor([0, 0, th.pi / 2])) + robot.set_position_orientation(position=new_world_pos, orientation=new_world_ori) - got_scene_pos, got_scene_ori = robot.get_position_orientation(frame="scene") - assert th.allclose(got_scene_pos, new_scene_pos, atol=1e-3) - assert th.allclose(got_scene_ori, new_scene_ori, atol=1e-3) + got_world_pos, got_world_ori = robot.get_position_orientation() + assert th.allclose(got_world_pos, new_world_pos, atol=1e-3) + assert th.allclose(got_world_ori, new_world_ori, atol=1e-3) - # Test setting position and orientation in scene frame - new_scene_pos = th.tensor([-1.0, -2.0, 0.1]) - new_scene_ori = T.euler2quat(th.tensor([th.pi / 6, 0, 0])) - robot.set_position_orientation(position=new_scene_pos, orientation=new_scene_ori, frame="scene") + # Test setting position and orientation in scene frame + new_scene_pos = th.tensor([0.5, 1.0, 0.25]) + new_scene_ori = T.euler2quat(th.tensor([0, th.pi / 4, 0])) + robot.set_position_orientation(position=new_scene_pos, orientation=new_scene_ori, frame="scene") - # Verify that world frame position/orientation has changed after setting in scene frame - got_world_pos, got_world_ori = robot.get_position_orientation() - assert not th.allclose(got_world_pos, new_world_pos, atol=1e-3) - assert not th.allclose(got_world_ori, new_world_ori, atol=1e-3) + got_scene_pos, got_scene_ori = robot.get_position_orientation(frame="scene") + assert th.allclose(got_scene_pos, new_scene_pos, atol=1e-3) + assert th.allclose(got_scene_ori, new_scene_ori, atol=1e-3) - # Clean up - og.clear() + # Test setting position and orientation in scene frame + new_scene_pos = th.tensor([-1.0, -2.0, 0.1]) + new_scene_ori = T.euler2quat(th.tensor([th.pi / 6, 0, 0])) + robot.set_position_orientation(position=new_scene_pos, orientation=new_scene_ori, frame="scene") - # assert that when the simulator is stopped, the behavior for getter/setter is not affected - vec_env = setup_multi_environment(2) - og.sim.stop() + # Verify that world frame position/orientation has changed after setting in scene frame + got_world_pos, got_world_ori = robot.get_position_orientation() + assert not th.allclose(got_world_pos, new_world_pos, atol=1e-3) + assert not th.allclose(got_world_ori, new_world_ori, atol=1e-3) - # use a robot with non-zero scene position - robot = vec_env.envs[1].scene.robots[0] + # Clean up + og.clear() - # Test setting position and orientation in world frame - new_world_pos = th.tensor([1.0, 2.0, 0.5]) - new_world_ori = T.euler2quat(th.tensor([0, 0, th.pi / 2])) - robot.set_position_orientation(position=new_world_pos, orientation=new_world_ori) - - got_world_pos, got_world_ori = robot.get_position_orientation() - assert th.allclose(got_world_pos, new_world_pos, atol=1e-3) - assert th.allclose(got_world_ori, new_world_ori, atol=1e-3) - - # Test setting position and orientation in scene frame - new_scene_pos = th.tensor([0.5, 1.0, 0.25]) - new_scene_ori = T.euler2quat(th.tensor([0, th.pi / 4, 0])) - robot.set_position_orientation(position=new_scene_pos, orientation=new_scene_ori, frame="scene") - - got_scene_pos, got_scene_ori = robot.get_position_orientation(frame="scene") - assert th.allclose(got_scene_pos, new_scene_pos, atol=1e-3) - assert th.allclose(got_scene_ori, new_scene_ori, atol=1e-3) - - # Test setting position and orientation in scene frame - new_scene_pos = th.tensor([-1.0, -2.0, 0.1]) - new_scene_ori = T.euler2quat(th.tensor([th.pi / 6, 0, 0])) - robot.set_position_orientation(position=new_scene_pos, orientation=new_scene_ori, frame="scene") - - got_scene_pos, got_scene_ori = robot.get_position_orientation(frame="scene") - assert th.allclose(got_scene_pos, new_scene_pos, atol=1e-3) - assert th.allclose(got_scene_ori, new_scene_ori, atol=1e-3) - - # Verify that world frame position/orientation has changed after setting in scene frame - got_world_pos, got_world_ori = robot.get_position_orientation() - assert not th.allclose(got_world_pos, new_world_pos, atol=1e-3) - assert not th.allclose(got_world_ori, new_world_ori, atol=1e-3) - - og.clear() - - -@pytest.mark.skip("Behavior getter is currently broken") -def test_behavior_getter(): - vec_env = setup_multi_environment(2, robot="BehaviorRobot") - robot1 = vec_env.envs[0].scene.robots[0] - - robot1_world_position, robot1_world_orientation = robot1.get_position_orientation() - robot1_scene_position, robot1_scene_orientation = robot1.get_position_orientation(frame="scene") + # assert that when the simulator is stopped, the behavior for getter/setter is not affected + vec_env = self.setup_multi_environment(pipeline_mode, 2) + og.sim.stop() - # Test the get_position_orientation method for 3 different frames - # since the robot is at the origin, the position and orientation should be the same - assert th.allclose(robot1_world_position, robot1_scene_position, atol=1e-3) - assert th.allclose(robot1_world_position, robot1_scene_position, atol=1e-3) - assert th.allclose(robot1_world_orientation, robot1_scene_orientation, atol=1e-3) - assert th.allclose(robot1_world_orientation, robot1_scene_orientation, atol=1e-3) + # use a robot with non-zero scene position + robot = vec_env.envs[1].scene.robots[0] + + # Test setting position and orientation in world frame + new_world_pos = th.tensor([1.0, 2.0, 0.5]) + new_world_ori = T.euler2quat(th.tensor([0, 0, th.pi / 2])) + robot.set_position_orientation(position=new_world_pos, orientation=new_world_ori) + + got_world_pos, got_world_ori = robot.get_position_orientation() + assert th.allclose(got_world_pos, new_world_pos, atol=1e-3) + assert th.allclose(got_world_ori, new_world_ori, atol=1e-3) + + # Test setting position and orientation in scene frame + new_scene_pos = th.tensor([0.5, 1.0, 0.25]) + new_scene_ori = T.euler2quat(th.tensor([0, th.pi / 4, 0])) + robot.set_position_orientation(position=new_scene_pos, orientation=new_scene_ori, frame="scene") + + got_scene_pos, got_scene_ori = robot.get_position_orientation(frame="scene") + assert th.allclose(got_scene_pos, new_scene_pos, atol=1e-3) + assert th.allclose(got_scene_ori, new_scene_ori, atol=1e-3) + + # Test setting position and orientation in scene frame + new_scene_pos = th.tensor([-1.0, -2.0, 0.1]) + new_scene_ori = T.euler2quat(th.tensor([th.pi / 6, 0, 0])) + robot.set_position_orientation(position=new_scene_pos, orientation=new_scene_ori, frame="scene") + + got_scene_pos, got_scene_ori = robot.get_position_orientation(frame="scene") + assert th.allclose(got_scene_pos, new_scene_pos, atol=1e-3) + assert th.allclose(got_scene_ori, new_scene_ori, atol=1e-3) + + # Verify that world frame position/orientation has changed after setting in scene frame + got_world_pos, got_world_ori = robot.get_position_orientation() + assert not th.allclose(got_world_pos, new_world_pos, atol=1e-3) + assert not th.allclose(got_world_ori, new_world_ori, atol=1e-3) + + og.clear() + + @pytest.mark.skip("Behavior getter is currently broken") + def test_behavior_getter(self, pipeline_mode): + vec_env = self.setup_multi_environment(pipeline_mode, 2, robot="BehaviorRobot") + robot1 = vec_env.envs[0].scene.robots[0] + + robot1_world_position, robot1_world_orientation = robot1.get_position_orientation() + robot1_scene_position, robot1_scene_orientation = robot1.get_position_orientation(frame="scene") + + # Test the get_position_orientation method for 3 different frames + # since the robot is at the origin, the position and orientation should be the same + assert th.allclose(robot1_world_position, robot1_scene_position, atol=1e-3) + assert th.allclose(robot1_world_position, robot1_scene_position, atol=1e-3) + assert th.allclose(robot1_world_orientation, robot1_scene_orientation, atol=1e-3) + assert th.allclose(robot1_world_orientation, robot1_scene_orientation, atol=1e-3) + + # test if the scene position is non-zero, the getter with scene and world frame should return different values + robot2 = vec_env.envs[1].scene.robots[0] + scene_position, scene_orientation = vec_env.envs[1].scene.get_position_orientation() + robot2_world_position, robot2_world_orientation = robot2.get_position_orientation() + robot2_scene_position, robot2_scene_orientation = robot2.get_position_orientation(frame="scene") + + combined_position, combined_orientation = T.pose_transform( + scene_position, scene_orientation, robot2_scene_position, robot2_scene_orientation + ) + assert th.allclose(robot2_world_position, combined_position, atol=1e-3) + assert th.allclose(robot2_world_orientation, combined_orientation, atol=1e-3) + + # Clean up + og.clear() + + @pytest.mark.skip("Behavior setter is currently broken") + def test_behavior_setter(self, pipeline_mode): + vec_env = self.setup_multi_environment(pipeline_mode, 2, robot="BehaviorRobot") + + # use a robot with non-zero scene position + robot = vec_env.envs[1].scene.robots[0] + + # Test setting position and orientation in world frame + new_world_pos = th.tensor([1.0, 2.0, 0.5]) + new_world_ori = T.euler2quat(th.tensor([0, 0, th.pi / 2])) + + robot.set_position_orientation(position=new_world_pos, orientation=new_world_ori) + + got_world_pos, got_world_ori = robot.get_position_orientation() + assert th.allclose(got_world_pos, new_world_pos, atol=1e-3) + assert th.allclose(got_world_ori, new_world_ori, atol=1e-3) + + # Test setting position and orientation in scene frame + new_scene_pos = th.tensor([0.5, 1.0, 0.25]) + new_scene_ori = T.euler2quat(th.tensor([0, th.pi / 4, 0])) + robot.set_position_orientation(position=new_scene_pos, orientation=new_scene_ori, frame="scene") + + got_scene_pos, got_scene_ori = robot.get_position_orientation(frame="scene") + assert th.allclose(got_scene_pos, new_scene_pos, atol=1e-3) + assert th.allclose(got_scene_ori, new_scene_ori, atol=1e-3) + + # Test setting position and orientation in scene frame + new_scene_pos = th.tensor([-1.0, -2.0, 0.1]) + new_scene_ori = T.euler2quat(th.tensor([th.pi / 6, 0, 0])) + robot.set_position_orientation(position=new_scene_pos, orientation=new_scene_ori, frame="scene") + + got_scene_pos, got_scene_ori = robot.get_position_orientation(frame="scene") + assert th.allclose(got_scene_pos, new_scene_pos, atol=1e-3) + assert th.allclose(got_scene_ori, new_scene_ori, atol=1e-3) + + # Verify that world frame position/orientation has changed after setting in scene frame + got_world_pos, got_world_ori = robot.get_position_orientation() + assert not th.allclose(got_world_pos, new_world_pos, atol=1e-3) + assert not th.allclose(got_world_ori, new_world_ori, atol=1e-3) + + # Clean up + og.clear() + + # assert that when the simulator is stopped, the behavior for getter/setter is not affected + vec_env = self.setup_multi_environment(pipeline_mode, 2) + og.sim.stop() - # test if the scene position is non-zero, the getter with scene and world frame should return different values - robot2 = vec_env.envs[1].scene.robots[0] - scene_position, scene_orientation = vec_env.envs[1].scene.get_position_orientation() - robot2_world_position, robot2_world_orientation = robot2.get_position_orientation() - robot2_scene_position, robot2_scene_orientation = robot2.get_position_orientation(frame="scene") - - combined_position, combined_orientation = T.pose_transform( - scene_position, scene_orientation, robot2_scene_position, robot2_scene_orientation - ) - assert th.allclose(robot2_world_position, combined_position, atol=1e-3) - assert th.allclose(robot2_world_orientation, combined_orientation, atol=1e-3) - - # Clean up - og.clear() - - -@pytest.mark.skip("Behavior setter is currently broken") -def test_behavior_setter(): - vec_env = setup_multi_environment(2, robot="BehaviorRobot") - - # use a robot with non-zero scene position - robot = vec_env.envs[1].scene.robots[0] - - # Test setting position and orientation in world frame - new_world_pos = th.tensor([1.0, 2.0, 0.5]) - new_world_ori = T.euler2quat(th.tensor([0, 0, th.pi / 2])) - - robot.set_position_orientation(position=new_world_pos, orientation=new_world_ori) - - got_world_pos, got_world_ori = robot.get_position_orientation() - assert th.allclose(got_world_pos, new_world_pos, atol=1e-3) - assert th.allclose(got_world_ori, new_world_ori, atol=1e-3) - - # Test setting position and orientation in scene frame - new_scene_pos = th.tensor([0.5, 1.0, 0.25]) - new_scene_ori = T.euler2quat(th.tensor([0, th.pi / 4, 0])) - robot.set_position_orientation(position=new_scene_pos, orientation=new_scene_ori, frame="scene") - - got_scene_pos, got_scene_ori = robot.get_position_orientation(frame="scene") - assert th.allclose(got_scene_pos, new_scene_pos, atol=1e-3) - assert th.allclose(got_scene_ori, new_scene_ori, atol=1e-3) - - # Test setting position and orientation in scene frame - new_scene_pos = th.tensor([-1.0, -2.0, 0.1]) - new_scene_ori = T.euler2quat(th.tensor([th.pi / 6, 0, 0])) - robot.set_position_orientation(position=new_scene_pos, orientation=new_scene_ori, frame="scene") + # use a robot with non-zero scene position + robot = vec_env.envs[1].scene.robots[0] - got_scene_pos, got_scene_ori = robot.get_position_orientation(frame="scene") - assert th.allclose(got_scene_pos, new_scene_pos, atol=1e-3) - assert th.allclose(got_scene_ori, new_scene_ori, atol=1e-3) - - # Verify that world frame position/orientation has changed after setting in scene frame - got_world_pos, got_world_ori = robot.get_position_orientation() - assert not th.allclose(got_world_pos, new_world_pos, atol=1e-3) - assert not th.allclose(got_world_ori, new_world_ori, atol=1e-3) - - # Clean up - og.clear() + # Test setting position and orientation in world frame + new_world_pos = th.tensor([1.0, 2.0, 0.5]) + new_world_ori = T.euler2quat(th.tensor([0, 0, th.pi / 2])) + robot.set_position_orientation(position=new_world_pos, orientation=new_world_ori) - # assert that when the simulator is stopped, the behavior for getter/setter is not affected - vec_env = setup_multi_environment(2) - og.sim.stop() + got_world_pos, got_world_ori = robot.get_position_orientation() + assert th.allclose(got_world_pos, new_world_pos, atol=1e-3) + assert th.allclose(got_world_ori, new_world_ori, atol=1e-3) - # use a robot with non-zero scene position - robot = vec_env.envs[1].scene.robots[0] - - # Test setting position and orientation in world frame - new_world_pos = th.tensor([1.0, 2.0, 0.5]) - new_world_ori = T.euler2quat(th.tensor([0, 0, th.pi / 2])) - robot.set_position_orientation(position=new_world_pos, orientation=new_world_ori) + # Test setting position and orientation in scene frame + new_scene_pos = th.tensor([0.5, 1.0, 0.25]) + new_scene_ori = T.euler2quat(th.tensor([0, th.pi / 4, 0])) + robot.set_position_orientation(position=new_scene_pos, orientation=new_scene_ori, frame="scene") - got_world_pos, got_world_ori = robot.get_position_orientation() - assert th.allclose(got_world_pos, new_world_pos, atol=1e-3) - assert th.allclose(got_world_ori, new_world_ori, atol=1e-3) + got_scene_pos, got_scene_ori = robot.get_position_orientation(frame="scene") + assert th.allclose(got_scene_pos, new_scene_pos, atol=1e-3) + assert th.allclose(got_scene_ori, new_scene_ori, atol=1e-3) - # Test setting position and orientation in scene frame - new_scene_pos = th.tensor([0.5, 1.0, 0.25]) - new_scene_ori = T.euler2quat(th.tensor([0, th.pi / 4, 0])) - robot.set_position_orientation(position=new_scene_pos, orientation=new_scene_ori, frame="scene") + # Test setting position and orientation in scene frame + new_scene_pos = th.tensor([-1.0, -2.0, 0.1]) + new_scene_ori = T.euler2quat(th.tensor([th.pi / 6, 0, 0])) + robot.set_position_orientation(position=new_scene_pos, orientation=new_scene_ori, frame="scene") - got_scene_pos, got_scene_ori = robot.get_position_orientation(frame="scene") - assert th.allclose(got_scene_pos, new_scene_pos, atol=1e-3) - assert th.allclose(got_scene_ori, new_scene_ori, atol=1e-3) + got_scene_pos, got_scene_ori = robot.get_position_orientation(frame="scene") + assert th.allclose(got_scene_pos, new_scene_pos, atol=1e-3) + assert th.allclose(got_scene_ori, new_scene_ori, atol=1e-3) - # Test setting position and orientation in scene frame - new_scene_pos = th.tensor([-1.0, -2.0, 0.1]) - new_scene_ori = T.euler2quat(th.tensor([th.pi / 6, 0, 0])) - robot.set_position_orientation(position=new_scene_pos, orientation=new_scene_ori, frame="scene") - - got_scene_pos, got_scene_ori = robot.get_position_orientation(frame="scene") - assert th.allclose(got_scene_pos, new_scene_pos, atol=1e-3) - assert th.allclose(got_scene_ori, new_scene_ori, atol=1e-3) - - # Verify that world frame position/orientation has changed after setting in scene frame - got_world_pos, got_world_ori = robot.get_position_orientation() - assert not th.allclose(got_world_pos, new_world_pos, atol=1e-3) - assert not th.allclose(got_world_ori, new_world_ori, atol=1e-3) + # Verify that world frame position/orientation has changed after setting in scene frame + got_world_pos, got_world_ori = robot.get_position_orientation() + assert not th.allclose(got_world_pos, new_world_pos, atol=1e-3) + assert not th.allclose(got_world_ori, new_world_ori, atol=1e-3) diff --git a/tests/test_object_removal.py b/tests/test_object_removal.py index bce9c889f..9bb75aacb 100644 --- a/tests/test_object_removal.py +++ b/tests/test_object_removal.py @@ -7,7 +7,7 @@ @og_test -def test_removal_and_readdition(env): +def test_removal_and_readdition(env, pipeline_mode): # Add an apple apple = DatasetObject( name="apple_unique", @@ -45,7 +45,7 @@ def test_removal_and_readdition(env): @og_test -def test_readdition(env): +def test_readdition(env, pipeline_mode): # Add an apple apple = DatasetObject( name="apple_unique", diff --git a/tests/test_object_states.py b/tests/test_object_states.py index 317a9dcc3..4acce7a54 100644 --- a/tests/test_object_states.py +++ b/tests/test_object_states.py @@ -6,6 +6,7 @@ import omnigibson as og import omnigibson.utils.transform_utils as T +from omnigibson.action_primitives.starter_semantic_action_primitives import StarterSemanticActionPrimitives from omnigibson.macros import macros as m from omnigibson.object_states import * from omnigibson.systems import VisualParticleSystem @@ -14,7 +15,7 @@ @og_test -def test_on_top(env): +def test_on_top(env, pipeline_mode): breakfast_table = env.scene.object_registry("name", "breakfast_table") bowl = env.scene.object_registry("name", "bowl") dishtowel = env.scene.object_registry("name", "dishtowel") @@ -31,16 +32,14 @@ def test_on_top(env): og.sim.step() assert not obj.states[OnTop].get_value(breakfast_table) - - assert bowl.states[OnTop].set_value(breakfast_table, True) - assert dishtowel.states[OnTop].set_value(breakfast_table, True) + assert obj.states[OnTop].set_value(breakfast_table, True) with pytest.raises(NotImplementedError): bowl.states[OnTop].set_value(breakfast_table, False) @og_test -def test_inside(env): +def test_inside(env, pipeline_mode): bottom_cabinet = env.scene.object_registry("name", "bottom_cabinet") bowl = env.scene.object_registry("name", "bowl") dishtowel = env.scene.object_registry("name", "dishtowel") @@ -72,7 +71,7 @@ def test_inside(env): @og_test -def test_under(env): +def test_under(env, pipeline_mode): breakfast_table = env.scene.object_registry("name", "breakfast_table") bowl = env.scene.object_registry("name", "bowl") dishtowel = env.scene.object_registry("name", "dishtowel") @@ -98,7 +97,7 @@ def test_under(env): @og_test -def test_touching(env): +def test_touching(env, pipeline_mode): breakfast_table = env.scene.object_registry("name", "breakfast_table") bowl = env.scene.object_registry("name", "bowl") dishtowel = env.scene.object_registry("name", "dishtowel") @@ -123,7 +122,7 @@ def test_touching(env): @og_test -def test_contact_bodies(env): +def test_contact_bodies(env, pipeline_mode): breakfast_table = env.scene.object_registry("name", "breakfast_table") bowl = env.scene.object_registry("name", "bowl") dishtowel = env.scene.object_registry("name", "dishtowel") @@ -150,7 +149,7 @@ def test_contact_bodies(env): @og_test -def test_next_to(env): +def test_next_to(env, pipeline_mode): bottom_cabinet = env.scene.object_registry("name", "bottom_cabinet") bowl = env.scene.object_registry("name", "bowl") dishtowel = env.scene.object_registry("name", "dishtowel") @@ -175,7 +174,9 @@ def test_next_to(env): @og_test -def test_overlaid(env): +def test_overlaid(env, pipeline_mode): + if pipeline_mode == "cpu": + pytest.skip("Overlaid requires cloth simulation, which is not supported in CPU mode") breakfast_table = env.scene.object_registry("name", "breakfast_table") carpet = env.scene.object_registry("name", "carpet") @@ -199,7 +200,7 @@ def test_overlaid(env): @og_test -def test_pose(env): +def test_pose(env, pipeline_mode): breakfast_table = env.scene.object_registry("name", "breakfast_table") dishtowel = env.scene.object_registry("name", "dishtowel") @@ -223,7 +224,7 @@ def test_pose(env): @og_test -def test_joint(env): +def test_joint(env, pipeline_mode): breakfast_table = env.scene.object_registry("name", "breakfast_table") bottom_cabinet = env.scene.object_registry("name", "bottom_cabinet") @@ -240,7 +241,7 @@ def test_joint(env): @og_test -def test_aabb(env): +def test_aabb(env, pipeline_mode): breakfast_table = env.scene.object_registry("name", "breakfast_table") dishtowel = env.scene.object_registry("name", "dishtowel") @@ -259,21 +260,22 @@ def test_aabb(env): (breakfast_table.states[AABB].get_value()[0] < pos1) & (pos1 < breakfast_table.states[AABB].get_value()[1]) ) - pp = dishtowel.root_link.compute_particle_positions() - offset = dishtowel.root_link.cloth_system.particle_contact_offset - particle_aabb = (pp.min(dim=0).values - offset, pp.max(dim=0).values + offset) - assert th.allclose(dishtowel.states[AABB].get_value()[0], particle_aabb[0]) - assert th.allclose(dishtowel.states[AABB].get_value()[1], particle_aabb[1]) - assert th.all( - (dishtowel.states[AABB].get_value()[0] < pos2) & (pos2 < dishtowel.states[AABB].get_value()[1]) - ).item() + if pipeline_mode == "cuda": + pp = dishtowel.root_link.compute_particle_positions() + offset = dishtowel.root_link.cloth_system.particle_contact_offset + particle_aabb = (pp.min(dim=0).values - offset, pp.max(dim=0).values + offset) + assert th.allclose(dishtowel.states[AABB].get_value()[0], particle_aabb[0]) + assert th.allclose(dishtowel.states[AABB].get_value()[1], particle_aabb[1]) + assert th.all( + (dishtowel.states[AABB].get_value()[0] < pos2) & (pos2 < dishtowel.states[AABB].get_value()[1]) + ).item() with pytest.raises(NotImplementedError): breakfast_table.states[AABB].set_value(None) @og_test -def test_adjacency(env): +def test_adjacency(env, pipeline_mode): bottom_cabinet = env.scene.object_registry("name", "bottom_cabinet") bowl = env.scene.object_registry("name", "bowl") dishtowel = env.scene.object_registry("name", "dishtowel") @@ -312,7 +314,7 @@ def test_adjacency(env): @og_test -def test_temperature(env): +def test_temperature(env, pipeline_mode): microwave = env.scene.object_registry("name", "microwave") stove = env.scene.object_registry("name", "stove") fridge = env.scene.object_registry("name", "fridge") @@ -446,7 +448,7 @@ def test_temperature(env): @og_test -def test_max_temperature(env): +def test_max_temperature(env, pipeline_mode): bagel = env.scene.object_registry("name", "bagel") dishtowel = env.scene.object_registry("name", "cookable_dishtowel") @@ -468,7 +470,7 @@ def test_max_temperature(env): @og_test -def test_heat_source_or_sink(env): +def test_heat_source_or_sink(env, pipeline_mode): microwave = env.scene.object_registry("name", "microwave") stove = env.scene.object_registry("name", "stove") fridge = env.scene.object_registry("name", "fridge") @@ -519,7 +521,7 @@ def test_heat_source_or_sink(env): @og_test -def test_cooked(env): +def test_cooked(env, pipeline_mode): bagel = env.scene.object_registry("name", "bagel") dishtowel = env.scene.object_registry("name", "cookable_dishtowel") @@ -548,7 +550,7 @@ def test_cooked(env): @og_test -def test_burnt(env): +def test_burnt(env, pipeline_mode): bagel = env.scene.object_registry("name", "bagel") dishtowel = env.scene.object_registry("name", "cookable_dishtowel") @@ -577,7 +579,7 @@ def test_burnt(env): @og_test -def test_frozen(env): +def test_frozen(env, pipeline_mode): bagel = env.scene.object_registry("name", "bagel") dishtowel = env.scene.object_registry("name", "cookable_dishtowel") @@ -606,7 +608,7 @@ def test_frozen(env): @og_test -def test_heated(env): +def test_heated(env, pipeline_mode): bagel = env.scene.object_registry("name", "bagel") dishtowel = env.scene.object_registry("name", "cookable_dishtowel") @@ -635,7 +637,7 @@ def test_heated(env): @og_test -def test_on_fire(env): +def test_on_fire(env, pipeline_mode): plywood = env.scene.object_registry("name", "plywood") assert not plywood.states[OnFire].get_value() @@ -659,32 +661,19 @@ def test_on_fire(env): assert plywood.states[Temperature].get_value() == plywood.states[OnFire].temperature +@pytest.mark.skip(reason="skipping toggle on test until trigger mesh is implemented") @og_test -def test_toggled_on(env): +def test_toggled_on(env, pipeline_mode): stove = env.scene.object_registry("name", "stove") robot = env.robots[0] - stove.set_position_orientation( - [1.48, 0.3, 0.443], T.euler2quat(th.tensor([0, 0, -math.pi / 2.0], dtype=th.float32)) - ) robot.set_position_orientation(position=[0.0, 0.38, 0.0], orientation=[0, 0, 0, 1]) + reset_joint_pos = robot.reset_joint_pos + robot.set_joint_positions(reset_joint_pos, drive=False) + stove.set_position_orientation([1.3, 0.3, 0.443], T.euler2quat(th.tensor([0, 0, -math.pi / 2.0], dtype=th.float32))) assert not stove.states[ToggledOn].get_value() - q = robot.get_joint_positions() - jnt_idxs = {name: i for i, name in enumerate(robot.joints.keys())} - q[jnt_idxs["torso_lift_joint"]] = 0.0 - q[jnt_idxs["shoulder_pan_joint"]] = th.deg2rad(th.tensor([90.0])).item() - q[jnt_idxs["shoulder_lift_joint"]] = th.deg2rad(th.tensor([9.0])).item() - q[jnt_idxs["upperarm_roll_joint"]] = 0.0 - q[jnt_idxs["elbow_flex_joint"]] = 0.0 - q[jnt_idxs["forearm_roll_joint"]] = 0.0 - q[jnt_idxs["wrist_flex_joint"]] = 0.0 - q[jnt_idxs["wrist_roll_joint"]] = 0.0 - q[jnt_idxs["l_gripper_finger_joint"]] = 0.0 - q[jnt_idxs["r_gripper_finger_joint"]] = 0.0 - robot.set_joint_positions(q, drive=False) - steps = m.object_states.toggle.CAN_TOGGLE_STEPS for _ in range(steps): og.sim.step() @@ -692,16 +681,26 @@ def test_toggled_on(env): # End-effector not close to the button, stays False assert not stove.states[ToggledOn].get_value() - # Settle robot - for _ in range(10): - og.sim.step() + # load IK controller + controller_config = { + f"arm_{robot.default_arm}": {"name": "InverseKinematicsController", "mode": "pose_absolute_ori"} + } + robot.reload_controllers(controller_config=controller_config) + + action_primitives = StarterSemanticActionPrimitives(env) + + target_eef_pos = stove.links["togglebutton_fifth_0_link"].get_position() + th.tensor( + [-0.02, 0.06, 0.0], dtype=th.float32 + ) + target_orn = T.quat_multiply( + robot.get_eef_orientation(), T.euler2quat(th.tensor([-math.pi / 2.0, 0.0, math.pi / 4.0])) + ) - q[jnt_idxs["shoulder_pan_joint"]] = 0.0 - robot.set_joint_positions(q, drive=False) + for action in action_primitives._move_hand_direct_ik((target_eef_pos, target_orn), pos_thresh=0.035): + env.step(action) for _ in range(steps - 1): og.sim.step() - robot.keep_still() # End-effector close to the button, but not enough time has passed, still False assert not stove.states[ToggledOn].get_value() @@ -717,7 +716,7 @@ def test_toggled_on(env): @og_test -def test_attached_to(env): +def test_attached_to(env, pipeline_mode): shelf_back_panel = env.scene.object_registry("name", "shelf_back_panel") shelf_shelf = env.scene.object_registry("name", "shelf_shelf") shelf_baseboard = env.scene.object_registry("name", "shelf_baseboard") @@ -788,7 +787,7 @@ def test_attached_to(env): @og_test -def test_particle_source(env): +def test_particle_source(env, pipeline_mode): sink = env.scene.object_registry("name", "sink") place_obj_on_floor_plane(sink) @@ -816,7 +815,7 @@ def test_particle_source(env): @og_test -def test_particle_sink(env): +def test_particle_sink(env, pipeline_mode): sink = env.scene.object_registry("name", "sink") place_obj_on_floor_plane(sink) for _ in range(3): @@ -845,7 +844,7 @@ def test_particle_sink(env): @og_test -def test_particle_applier(env): +def test_particle_applier(env, pipeline_mode): breakfast_table = env.scene.object_registry("name", "breakfast_table") spray_bottle = env.scene.object_registry("name", "spray_bottle") applier_dishtowel = env.scene.object_registry("name", "applier_dishtowel") @@ -906,7 +905,7 @@ def test_particle_applier(env): @og_test -def test_particle_remover(env): +def test_particle_remover(env, pipeline_mode): breakfast_table = env.scene.object_registry("name", "breakfast_table") vacuum = env.scene.object_registry("name", "vacuum") remover_dishtowel = env.scene.object_registry("name", "remover_dishtowel") @@ -971,7 +970,7 @@ def test_particle_remover(env): @og_test -def test_saturated(env): +def test_saturated(env, pipeline_mode): remover_dishtowel = env.scene.object_registry("name", "remover_dishtowel") place_obj_on_floor_plane(remover_dishtowel) @@ -1007,7 +1006,7 @@ def test_saturated(env): @og_test -def test_open(env): +def test_open(env, pipeline_mode): microwave = env.scene.object_registry("name", "microwave") bottom_cabinet = env.scene.object_registry("name", "bottom_cabinet") @@ -1053,7 +1052,7 @@ def test_open(env): @og_test -def test_folded_unfolded(env): +def test_folded_unfolded(env, pipeline_mode): carpet = env.scene.object_registry("name", "carpet") place_obj_on_floor_plane(carpet) @@ -1105,7 +1104,7 @@ def test_folded_unfolded(env): @og_test -def test_draped(env): +def test_draped(env, pipeline_mode): breakfast_table = env.scene.object_registry("name", "breakfast_table") carpet = env.scene.object_registry("name", "carpet") @@ -1131,7 +1130,7 @@ def test_draped(env): @og_test -def test_filled(env): +def test_filled(env, pipeline_mode): stockpot = env.scene.object_registry("name", "stockpot") systems = [ env.scene.get_system(system_name) @@ -1158,7 +1157,7 @@ def test_filled(env): @og_test -def test_contains(env): +def test_contains(env, pipeline_mode): stockpot = env.scene.object_registry("name", "stockpot") systems = [env.scene.get_system(system_name) for system_name, system_class in SYSTEM_EXAMPLES.items()] for system in systems: @@ -1196,7 +1195,7 @@ def test_contains(env): @og_test -def test_covered(env): +def test_covered(env, pipeline_mode): bracelet = env.scene.object_registry("name", "bracelet") bowl = env.scene.object_registry("name", "bowl") microwave = env.scene.object_registry("name", "microwave") diff --git a/tests/test_primitives.py b/tests/test_primitives.py index f77827c9d..2682a1aeb 100644 --- a/tests/test_primitives.py +++ b/tests/test_primitives.py @@ -52,8 +52,6 @@ def setup_environment(load_object_categories, robot="Fetch"): if og.sim is None: # Make sure GPU dynamics are enabled (GPU dynamics needed for cloth) and no flatcache gm.ENABLE_OBJECT_STATES = True - gm.USE_GPU_DYNAMICS = False - gm.ENABLE_FLATCACHE = False gm.ENABLE_TRANSITION_RULES = False else: # Make sure sim is stopped diff --git a/tests/test_robot_states.py b/tests/test_robot_states.py new file mode 100644 index 000000000..9b748cb92 --- /dev/null +++ b/tests/test_robot_states.py @@ -0,0 +1,244 @@ +import pytest +import torch as th + +import omnigibson as og +import omnigibson.lazy as lazy +from omnigibson.action_primitives.starter_semantic_action_primitives import StarterSemanticActionPrimitives +from omnigibson.macros import gm +from omnigibson.object_states import ObjectsInFOVOfRobot +from omnigibson.robots import * +from omnigibson.sensors import VisionSensor +from omnigibson.utils.constants import semantic_class_name_to_id +from omnigibson.utils.transform_utils import mat2pose, pose2mat, quaternions_close, relative_pose_transform +from omnigibson.utils.usd_utils import PoseAPI + + +@pytest.mark.parametrize("pipeline_mode", ["cpu", "cuda"], indirect=True) +class TestRobotStates: + def setup_environment(self, pipeline_mode): + """ + Sets up the environment. + """ + if og.sim is None: + # Set global flags + gm.ENABLE_OBJECT_STATES = True + gm.ENABLE_TRANSITION_RULES = False + else: + # Make sure sim is stopped + og.sim.stop() + + # Define the environment configuration + config = { + "env": { + "device": pipeline_mode, + }, + "scene": { + "type": "Scene", + }, + "robots": [ + { + "type": "Fetch", + "obs_modalities": ["rgb", "seg_semantic", "seg_instance"], + "position": [150, 150, 100], + "orientation": [0, 0, 0, 1], + } + ], + } + + env = og.Environment(configs=config) + return env + + def test_camera_pose(self, pipeline_mode): + env = self.setup_environment(pipeline_mode) + robot = env.robots[0] + env.reset() + og.sim.step() + + sensors = [s for s in robot.sensors.values() if isinstance(s, VisionSensor)] + assert len(sensors) > 0 + vision_sensor = sensors[0] + + # Get vision sensor world pose via directly calling get_position_orientation + robot_world_pos, robot_world_ori = robot.get_position_orientation() + sensor_world_pos, sensor_world_ori = vision_sensor.get_position_orientation() + + robot_to_sensor_mat = pose2mat( + relative_pose_transform(sensor_world_pos, sensor_world_ori, robot_world_pos, robot_world_ori) + ) + + sensor_world_pos_gt = th.tensor([150.1620, 149.9999, 101.2193]) + sensor_world_ori_gt = th.tensor([-0.2952, 0.2959, 0.6427, -0.6421]) + + assert th.allclose(sensor_world_pos, sensor_world_pos_gt, atol=1e-3) + assert quaternions_close(sensor_world_ori, sensor_world_ori_gt, atol=1e-3) + + # Now, we want to move the robot and check if the sensor pose has been updated + old_camera_local_pose = vision_sensor.get_position_orientation(frame="parent") + + robot.set_position_orientation(position=[100, 100, 100]) + new_camera_local_pose = vision_sensor.get_position_orientation(frame="parent") + new_camera_world_pose = vision_sensor.get_position_orientation() + robot_pose_mat = pose2mat(robot.get_position_orientation()) + expected_camera_world_pos, expected_camera_world_ori = mat2pose(robot_pose_mat @ robot_to_sensor_mat) + assert th.allclose(old_camera_local_pose[0], new_camera_local_pose[0], atol=1e-3) + assert th.allclose(new_camera_world_pose[0], expected_camera_world_pos, atol=1e-3) + assert quaternions_close(new_camera_world_pose[1], expected_camera_world_ori, atol=1e-3) + + # Then, we want to move the local pose of the camera and check + # 1) if the world pose is updated 2) if the robot stays in the same position + old_camera_local_pose = vision_sensor.get_position_orientation(frame="parent") + vision_sensor.set_position_orientation(position=[10, 10, 10], orientation=[0, 0, 0, 1], frame="parent") + new_camera_world_pose = vision_sensor.get_position_orientation() + camera_parent_prim = lazy.omni.isaac.core.utils.prims.get_prim_parent(vision_sensor.prim) + camera_parent_path = str(camera_parent_prim.GetPath()) + camera_parent_world_transform = PoseAPI.get_world_pose_with_scale(camera_parent_path) + expected_new_camera_world_pos, expected_new_camera_world_ori = mat2pose( + camera_parent_world_transform + @ pose2mat((th.tensor([10, 10, 10], dtype=th.float32), th.tensor([0, 0, 0, 1], dtype=th.float32))) + ) + assert th.allclose(new_camera_world_pose[0], expected_new_camera_world_pos, atol=1e-3) + assert quaternions_close(new_camera_world_pose[1], expected_new_camera_world_ori, atol=1e-3) + assert th.allclose(robot.get_position_orientation()[0], th.tensor([100, 100, 100], dtype=th.float32), atol=1e-3) + + # Finally, we want to move the world pose of the camera and check + # 1) if the local pose is updated 2) if the robot stays in the same position + robot.set_position_orientation(position=[150, 150, 100]) + old_camera_local_pose = vision_sensor.get_position_orientation(frame="parent") + vision_sensor.set_position_orientation( + position=[150, 150, 101.36912537], orientation=[-0.29444987, 0.29444981, 0.64288363, -0.64288352] + ) + new_camera_local_pose = vision_sensor.get_position_orientation(frame="parent") + assert not th.allclose(old_camera_local_pose[0], new_camera_local_pose[0], atol=1e-3) + assert not quaternions_close(old_camera_local_pose[1], new_camera_local_pose[1], atol=1e-3) + assert th.allclose(robot.get_position_orientation()[0], th.tensor([150, 150, 100], dtype=th.float32), atol=1e-3) + + # Another test we want to try is setting the camera's parent scale and check if the world pose is updated + camera_parent_prim.GetAttribute("xformOp:scale").Set(lazy.pxr.Gf.Vec3d([2.0, 2.0, 2.0])) + camera_parent_world_transform = PoseAPI.get_world_pose_with_scale(camera_parent_path) + camera_local_pose = vision_sensor.get_position_orientation(frame="parent") + expected_mat = camera_parent_world_transform @ pose2mat(camera_local_pose) + expected_mat[:3, :3] = expected_mat[:3, :3] / th.norm(expected_mat[:3, :3], dim=0, keepdim=True) + expected_new_camera_world_pos, _ = mat2pose(expected_mat) + new_camera_world_pose = vision_sensor.get_position_orientation() + assert th.allclose(new_camera_world_pose[0], expected_new_camera_world_pos, atol=1e-3) + + og.clear() + + def test_camera_semantic_segmentation(self, pipeline_mode): + env = self.setup_environment(pipeline_mode) + robot = env.robots[0] + env.reset() + sensors = [s for s in robot.sensors.values() if isinstance(s, VisionSensor)] + assert len(sensors) > 0 + vision_sensor = sensors[0] + env.reset() + all_observation, all_info = vision_sensor.get_obs() + seg_semantic = all_observation["seg_semantic"] + seg_semantic_info = all_info["seg_semantic"] + agent_label = semantic_class_name_to_id()["agent"] + background_label = semantic_class_name_to_id()["background"] + assert th.all(th.isin(seg_semantic, th.tensor([agent_label, background_label], device=seg_semantic.device))) + assert set(seg_semantic_info.keys()) == {agent_label, background_label} + og.clear() + + def test_object_in_FOV_of_robot(self, pipeline_mode): + env = self.setup_environment(pipeline_mode) + robot = env.robots[0] + env.reset() + assert robot.states[ObjectsInFOVOfRobot].get_value() == [robot] + sensors = [s for s in robot.sensors.values() if isinstance(s, VisionSensor)] + assert len(sensors) > 0 + vision_sensor = sensors[0] + vision_sensor.set_position_orientation(position=[100, 150, 100]) + og.sim.step() + og.sim.step() + assert robot.states[ObjectsInFOVOfRobot].get_value() == [robot] + og.clear() + + def test_robot_load_drive(self, pipeline_mode): + if og.sim is None: + # Set global flags + gm.ENABLE_OBJECT_STATES = True + gm.ENABLE_TRANSITION_RULES = False + else: + # Make sure sim is stopped + og.sim.stop() + + config = { + "env": { + "device": pipeline_mode, + }, + "scene": { + "type": "Scene", + }, + } + + env = og.Environment(configs=config) + og.sim.stop() + + # Iterate over all robots and test their motion + for robot_name, robot_cls in REGISTERED_ROBOTS.items(): + + if robot_name in ["FrankaMounted", "Stretch"]: + # TODO: skipping FrankaMounted and Stretch for now because CI doesn't have the required assets + continue + + robot = robot_cls( + name=robot_name, + obs_modalities=[], + ) + env.scene.add_object(robot) + + # At least one step is always needed while sim is playing for any imported object to be fully initialized + og.sim.play() + og.sim.step() + + # Reset robot and make sure it's not moving + robot.reset() + robot.keep_still() + + # Set viewer in front facing robot + og.sim.viewer_camera.set_position_orientation( + position=[2.69918369, -3.63686664, 4.57894564], + orientation=[0.39592411, 0.1348514, 0.29286304, 0.85982], + ) + + if not robot_name in ["Husky", "BehaviorRobot"]: + # Husky base motion is a little messed up because of the 4-wheel drive; skipping for now + # BehaviorRobot does not work with the primitive actions at the moment + + # If this is a manipulation robot, we want to test moving the arm + if isinstance(robot, ManipulationRobot): + # load IK controller + controller_config = { + f"arm_{robot.default_arm}": {"name": "InverseKinematicsController", "mode": "pose_absolute_ori"} + } + robot.reload_controllers(controller_config=controller_config) + env.scene.update_initial_state() + + action_primitives = StarterSemanticActionPrimitives(env) + + eef_pos = env.robots[0].get_eef_position() + eef_orn = env.robots[0].get_eef_orientation() + if isinstance(robot, Stretch): # Stretch arm faces the y-axis + target_eef_pos = th.tensor([eef_pos[0], eef_pos[1] - 0.1, eef_pos[2]], dtype=th.float32) + else: + target_eef_pos = th.tensor([eef_pos[0] + 0.1, eef_pos[1], eef_pos[2]], dtype=th.float32) + target_eef_orn = eef_orn + for action in action_primitives._move_hand_direct_ik((target_eef_pos, target_eef_orn)): + env.step(action) + assert th.norm(robot.get_eef_position() - target_eef_pos) < 0.05 + + # If this is a locomotion robot, we want to test driving + if isinstance(robot, LocomotionRobot): + action_primitives = StarterSemanticActionPrimitives(env) + goal_location = th.tensor([0, 1, 0], dtype=th.float32) + for action in action_primitives._navigate_to_pose_direct(goal_location): + env.step(action) + assert th.norm(robot.get_position_orientation()[0][:2] - goal_location[:2]) < 0.1 + + # Stop the simulator and remove the robot + og.sim.stop() + env.scene.remove_object(obj=robot) + + env.close() diff --git a/tests/test_robot_states_flatcache.py b/tests/test_robot_states_flatcache.py deleted file mode 100644 index 10c4a4c45..000000000 --- a/tests/test_robot_states_flatcache.py +++ /dev/null @@ -1,325 +0,0 @@ -import torch as th - -import omnigibson as og -import omnigibson.lazy as lazy -from omnigibson.action_primitives.starter_semantic_action_primitives import StarterSemanticActionPrimitives -from omnigibson.macros import gm -from omnigibson.robots import * -from omnigibson.sensors import VisionSensor -from omnigibson.utils.transform_utils import mat2pose, pose2mat, quaternions_close, relative_pose_transform -from omnigibson.utils.usd_utils import PoseAPI - - -def setup_environment(flatcache): - """ - Sets up the environment with or without flatcache based on the flatcache parameter. - """ - if og.sim is None: - # Set global flags - gm.ENABLE_OBJECT_STATES = True - gm.USE_GPU_DYNAMICS = True - gm.ENABLE_FLATCACHE = flatcache # Set based on function parameter - gm.ENABLE_TRANSITION_RULES = False - else: - # Make sure sim is stopped - og.sim.stop() - - # Define the environment configuration - config = { - "scene": { - "type": "Scene", - }, - "robots": [ - { - "type": "Fetch", - "obs_modalities": ["rgb", "seg_semantic", "seg_instance"], - "position": [150, 150, 100], - "orientation": [0, 0, 0, 1], - } - ], - } - - env = og.Environment(configs=config) - return env - - -def camera_pose_test(flatcache): - env = setup_environment(flatcache) - robot = env.robots[0] - env.reset() - og.sim.step() - - sensors = [s for s in robot.sensors.values() if isinstance(s, VisionSensor)] - assert len(sensors) > 0 - vision_sensor = sensors[0] - - # Get vision sensor world pose via directly calling get_position_orientation - robot_world_pos, robot_world_ori = robot.get_position_orientation() - sensor_world_pos, sensor_world_ori = vision_sensor.get_position_orientation() - - robot_to_sensor_mat = pose2mat( - relative_pose_transform(sensor_world_pos, sensor_world_ori, robot_world_pos, robot_world_ori) - ) - - sensor_world_pos_gt = th.tensor([150.1628, 149.9993, 101.3773]) - sensor_world_ori_gt = th.tensor([-0.2952, 0.2959, 0.6427, -0.6421]) - - assert th.allclose(sensor_world_pos, sensor_world_pos_gt, atol=1e-3) - assert quaternions_close(sensor_world_ori, sensor_world_ori_gt, atol=1e-3) - - # Now, we want to move the robot and check if the sensor pose has been updated - old_camera_local_pose = vision_sensor.get_position_orientation(frame="parent") - - robot.set_position_orientation(position=[100, 100, 100]) - new_camera_local_pose = vision_sensor.get_position_orientation(frame="parent") - new_camera_world_pose = vision_sensor.get_position_orientation() - robot_pose_mat = pose2mat(robot.get_position_orientation()) - expected_camera_world_pos, expected_camera_world_ori = mat2pose(robot_pose_mat @ robot_to_sensor_mat) - assert th.allclose(old_camera_local_pose[0], new_camera_local_pose[0], atol=1e-3) - assert th.allclose(new_camera_world_pose[0], expected_camera_world_pos, atol=1e-3) - assert quaternions_close(new_camera_world_pose[1], expected_camera_world_ori, atol=1e-3) - - # Then, we want to move the local pose of the camera and check - # 1) if the world pose is updated 2) if the robot stays in the same position - old_camera_local_pose = vision_sensor.get_position_orientation(frame="parent") - vision_sensor.set_position_orientation(position=[10, 10, 10], orientation=[0, 0, 0, 1], frame="parent") - new_camera_world_pose = vision_sensor.get_position_orientation() - camera_parent_prim = lazy.omni.isaac.core.utils.prims.get_prim_parent(vision_sensor.prim) - camera_parent_path = str(camera_parent_prim.GetPath()) - camera_parent_world_transform = PoseAPI.get_world_pose_with_scale(camera_parent_path) - expected_new_camera_world_pos, expected_new_camera_world_ori = mat2pose( - camera_parent_world_transform - @ pose2mat((th.tensor([10, 10, 10], dtype=th.float32), th.tensor([0, 0, 0, 1], dtype=th.float32))) - ) - assert th.allclose(new_camera_world_pose[0], expected_new_camera_world_pos, atol=1e-3) - assert quaternions_close(new_camera_world_pose[1], expected_new_camera_world_ori, atol=1e-3) - assert th.allclose(robot.get_position_orientation()[0], th.tensor([100, 100, 100], dtype=th.float32), atol=1e-3) - - # Finally, we want to move the world pose of the camera and check - # 1) if the local pose is updated 2) if the robot stays in the same position - robot.set_position_orientation(position=[150, 150, 100]) - old_camera_local_pose = vision_sensor.get_position_orientation(frame="parent") - vision_sensor.set_position_orientation( - position=[150, 150, 101.36912537], orientation=[-0.29444987, 0.29444981, 0.64288363, -0.64288352] - ) - new_camera_local_pose = vision_sensor.get_position_orientation(frame="parent") - assert not th.allclose(old_camera_local_pose[0], new_camera_local_pose[0], atol=1e-3) - assert not quaternions_close(old_camera_local_pose[1], new_camera_local_pose[1], atol=1e-3) - assert th.allclose(robot.get_position_orientation()[0], th.tensor([150, 150, 100], dtype=th.float32), atol=1e-3) - - # Another test we want to try is setting the camera's parent scale and check if the world pose is updated - camera_parent_prim.GetAttribute("xformOp:scale").Set(lazy.pxr.Gf.Vec3d([2.0, 2.0, 2.0])) - camera_parent_world_transform = PoseAPI.get_world_pose_with_scale(camera_parent_path) - camera_local_pose = vision_sensor.get_position_orientation(frame="parent") - expected_mat = camera_parent_world_transform @ pose2mat(camera_local_pose) - expected_mat[:3, :3] = expected_mat[:3, :3] / th.norm(expected_mat[:3, :3], dim=0, keepdim=True) - expected_new_camera_world_pos, _ = mat2pose(expected_mat) - new_camera_world_pose = vision_sensor.get_position_orientation() - assert th.allclose(new_camera_world_pose[0], expected_new_camera_world_pos, atol=1e-3) - - og.clear() - - -def test_camera_pose_flatcache_on(): - camera_pose_test(True) - - -def test_robot_load_drive(): - if og.sim is None: - # Set global flags - gm.ENABLE_OBJECT_STATES = True - gm.USE_GPU_DYNAMICS = True - gm.ENABLE_TRANSITION_RULES = False - else: - # Make sure sim is stopped - og.sim.stop() - - config = { - "scene": { - "type": "Scene", - }, - } - - env = og.Environment(configs=config) - og.sim.stop() - - # Iterate over all robots and test their motion - for robot_name, robot_cls in REGISTERED_ROBOTS.items(): - - if robot_name in ["FrankaMounted", "Stretch"]: - # TODO: skipping FrankaMounted and Stretch for now because CI doesn't have the required assets - continue - - if robot_name in ["Husky", "BehaviorRobot"]: - # Husky base motion is a little messed up because of the 4-wheel drive; skipping for now - # BehaviorRobot does not work with the primitive actions at the moment - continue - - robot = robot_cls( - name=robot_name, - obs_modalities=[], - ) - env.scene.add_object(robot) - - # At least one step is always needed while sim is playing for any imported object to be fully initialized - og.sim.play() - - # Reset robot and make sure it's not moving - robot.reset() - robot.keep_still() - - og.sim.step() - - # Set viewer in front facing robot - og.sim.viewer_camera.set_position_orientation( - position=[2.69918369, -3.63686664, 4.57894564], - orientation=[0.39592411, 0.1348514, 0.29286304, 0.85982], - ) - - # If this is a manipulation robot, we want to test moving the arm - if isinstance(robot, ManipulationRobot): - # load IK controller - controller_config = { - f"arm_{robot.default_arm}": {"name": "InverseKinematicsController", "mode": "pose_absolute_ori"} - } - robot.reload_controllers(controller_config=controller_config) - env.scene.update_initial_state() - - action_primitives = StarterSemanticActionPrimitives(env) - - eef_pos = env.robots[0].get_eef_position() - eef_orn = env.robots[0].get_eef_orientation() - if isinstance(robot, Stretch): # Stretch arm faces the y-axis - target_eef_pos = th.tensor([eef_pos[0], eef_pos[1] - 0.1, eef_pos[2]], dtype=th.float32) - else: - target_eef_pos = th.tensor([eef_pos[0] + 0.1, eef_pos[1], eef_pos[2]], dtype=th.float32) - target_eef_orn = eef_orn - for action in action_primitives._move_hand_direct_ik((target_eef_pos, target_eef_orn)): - env.step(action) - assert th.norm(robot.get_eef_position() - target_eef_pos) < 0.05 - - # If this is a locomotion robot, we want to test driving - if isinstance(robot, LocomotionRobot): - action_primitives = StarterSemanticActionPrimitives(env) - goal_location = th.tensor([0, 1, 0], dtype=th.float32) - for action in action_primitives._navigate_to_pose_direct(goal_location): - env.step(action) - assert th.norm(robot.get_position()[:2] - goal_location[:2]) < 0.1 - - # Stop the simulator and remove the robot - og.sim.stop() - env.scene.remove_object(obj=robot) - - og.clear() - - -def test_grasping_mode(): - - if og.sim is None: - # Set global flags - gm.ENABLE_FLATCACHE = True - else: - # Make sure sim is stopped - og.sim.stop() - - scene_cfg = dict(type="Scene") - objects_cfg = [] - objects_cfg.append( - dict( - type="DatasetObject", - name=f"table", - category="breakfast_table", - model="lcsizg", - bounding_box=[0.5, 0.5, 0.8], - fixed_base=True, - position=[0.7, -0.1, 0.6], - ) - ) - objects_cfg.append( - dict( - type="PrimitiveObject", - name=f"box", - primitive_type="Cube", - rgba=[1.0, 0, 0, 1.0], - size=0.05, - position=[0.53, 0.0, 0.87], - ) - ) - cfg = dict(scene=scene_cfg, objects=objects_cfg) - - env = og.Environment(configs=cfg) - og.sim.viewer_camera.set_position_orientation( - position=[1.0170, 0.5663, 1.0554], - orientation=[0.1734, 0.5006, 0.8015, 0.2776], - ) - og.sim.stop() - - grasping_modes = dict( - sticky="Sticky Mitten - Objects are magnetized when they touch the fingers and a CLOSE command is given", - assisted="Assisted Grasping - Objects are magnetized when they touch the fingers, are within the hand, and a CLOSE command is given", - physical="Physical Grasping - No additional grasping assistance applied", - ) - - def object_is_in_hand(robot, obj): - eef_position = robot.get_eef_position() - obj_position = obj.get_position_orientation()[0] - return th.norm(eef_position - obj_position) < 0.05 - - for grasping_mode in grasping_modes: - robot = Fetch( - name="Fetch", - obs_modalities=[], - controller_config={f"arm_0": {"name": "InverseKinematicsController", "mode": "pose_absolute_ori"}}, - grasping_mode=grasping_mode, - ) - env.scene.add_object(robot) - - # At least one step is always needed while sim is playing for any imported object to be fully initialized - og.sim.play() - - env.scene.reset() - - # Reset robot and make sure it's not moving - robot.reset() - robot.keep_still() - - # Let the box settle - for _ in range(10): - og.sim.step() - - action_primitives = StarterSemanticActionPrimitives(env) - - box_object = env.scene.object_registry("name", "box") - target_eef_pos = box_object.get_position_orientation()[0] - target_eef_orn = robot.get_eef_orientation() - - # Move eef to the box - for action in action_primitives._move_hand_direct_ik((target_eef_pos, target_eef_orn), pos_thresh=0.01): - env.step(action) - - # Grasp the box - for action in action_primitives._execute_grasp(): - env.step(action) - - assert object_is_in_hand(robot, box_object), f"Grasping mode {grasping_mode} failed to grasp the object" - - # Move eef - eef_offset = th.tensor([0.0, 0.2, 0.2]) - for action in action_primitives._move_hand_direct_ik((target_eef_pos + eef_offset, target_eef_orn)): - env.step(action) - - assert object_is_in_hand(robot, box_object), f"Grasping mode {grasping_mode} failed to keep the object in hand" - - # Release the box - for action in action_primitives._execute_release(): - env.step(action) - for _ in range(10): - og.sim.step() - - assert not object_is_in_hand(robot, box_object), f"Grasping mode {grasping_mode} failed to release the object" - - # Stop the simulator and remove the robot - og.sim.stop() - env.scene.remove_object(obj=robot) - - og.clear() diff --git a/tests/test_robot_states_no_flatcache.py b/tests/test_robot_states_no_flatcache.py deleted file mode 100644 index b02ba6158..000000000 --- a/tests/test_robot_states_no_flatcache.py +++ /dev/null @@ -1,45 +0,0 @@ -import torch as th -from test_robot_states_flatcache import camera_pose_test, setup_environment - -import omnigibson as og -from omnigibson.object_states import ObjectsInFOVOfRobot -from omnigibson.sensors import VisionSensor -from omnigibson.utils.constants import semantic_class_name_to_id - - -def test_camera_pose_flatcache_off(): - camera_pose_test(False) - - -def test_camera_semantic_segmentation(): - env = setup_environment(False) - robot = env.robots[0] - env.reset() - sensors = [s for s in robot.sensors.values() if isinstance(s, VisionSensor)] - assert len(sensors) > 0 - vision_sensor = sensors[0] - env.reset() - all_observation, all_info = vision_sensor.get_obs() - seg_semantic = all_observation["seg_semantic"] - seg_semantic_info = all_info["seg_semantic"] - agent_label = semantic_class_name_to_id()["agent"] - background_label = semantic_class_name_to_id()["background"] - assert th.all(th.isin(seg_semantic, th.tensor([agent_label, background_label], device=seg_semantic.device))) - assert set(seg_semantic_info.keys()) == {agent_label, background_label} - og.clear() - - -def test_object_in_FOV_of_robot(): - env = setup_environment(False) - robot = env.robots[0] - env.reset() - assert robot.states[ObjectsInFOVOfRobot].get_value() == [robot] - sensors = [s for s in robot.sensors.values() if isinstance(s, VisionSensor)] - assert len(sensors) > 0 - vision_sensor = sensors[0] - vision_sensor.set_position_orientation(position=[100, 150, 100]) - og.sim.step() - og.sim.step() - # Since the sensor is moved away from the robot, the robot should not see itself - assert robot.states[ObjectsInFOVOfRobot].get_value() == [] - og.clear() diff --git a/tests/test_robot_teleoperation.py b/tests/test_robot_teleoperation.py index 1e922b997..96f0adb43 100644 --- a/tests/test_robot_teleoperation.py +++ b/tests/test_robot_teleoperation.py @@ -8,9 +8,10 @@ @pytest.mark.skip(reason="test hangs on CI") -def test_teleop(): +@pytest.mark.parametrize("pipeline_mode", ["cpu", "cuda"], indirect=True) +def test_teleop(pipeline_mode): cfg = { - "env": {"action_timestep": 1 / 60.0, "physics_timestep": 1 / 120.0}, + "env": {"action_timestep": 1 / 60.0, "physics_timestep": 1 / 120.0, "device": pipeline_mode}, "scene": {"type": "Scene"}, "robots": [ { @@ -26,11 +27,7 @@ def test_teleop(): ], } - if og.sim is None: - # Make sure GPU dynamics are enabled (GPU dynamics needed for cloth) and no flatcache - gm.USE_GPU_DYNAMICS = False - gm.ENABLE_FLATCACHE = False - else: + if og.sim is not None: # Make sure sim is stopped og.sim.stop() diff --git a/tests/test_scene_graph.py b/tests/test_scene_graph.py index 50227a613..cecc45d54 100644 --- a/tests/test_scene_graph.py +++ b/tests/test_scene_graph.py @@ -1,5 +1,6 @@ import math +import pytest import torch as th from utils import place_obj_on_floor_plane @@ -10,7 +11,8 @@ from omnigibson.utils.constants import PrimType -def test_scene_graph(): +@pytest.mark.parametrize("pipeline_mode", ["cpu", "cuda"], indirect=True) +def test_scene_graph(pipeline_mode): if og.sim is None: # Set global flags @@ -38,6 +40,9 @@ def create_robot_config(name, position): robot_positions = [[0, 0.8, 0], [1, 0.8, 0], [2, 0.8, 0]] config = { + "env": { + "device": pipeline_mode, + }, "scene": { "type": "Scene", }, diff --git a/tests/test_sensors.py b/tests/test_sensors.py index b7d0ff32d..21f4f4102 100644 --- a/tests/test_sensors.py +++ b/tests/test_sensors.py @@ -10,7 +10,7 @@ @og_test -def test_segmentation_modalities(env): +def test_segmentation_modalities(env, pipeline_mode): breakfast_table = env.scene.object_registry("name", "breakfast_table") dishtowel = env.scene.object_registry("name", "dishtowel") robot = env.scene.robots[0] @@ -93,7 +93,7 @@ def test_segmentation_modalities(env): @og_test -def test_bbox_modalities(env): +def test_bbox_modalities(env, pipeline_mode): breakfast_table = env.scene.object_registry("name", "breakfast_table") dishtowel = env.scene.object_registry("name", "dishtowel") robot = env.scene.robots[0] diff --git a/tests/test_symbolic_primitives.py b/tests/test_symbolic_primitives.py index 4793686b7..d83ae0011 100644 --- a/tests/test_symbolic_primitives.py +++ b/tests/test_symbolic_primitives.py @@ -16,7 +16,6 @@ from omnigibson.robots import REGISTERED_ROBOTS from omnigibson.utils.python_utils import create_class_from_registry_and_config -gm.USE_GPU_DYNAMICS = True gm.ENABLE_TRANSITION_RULES = True current_robot_type = "Fetch" @@ -83,6 +82,9 @@ def start_env(robot_type): }, ], } + + gm.ENABLE_TRANSITION_RULES = False + env = og.Environment(configs=config) return env diff --git a/tests/test_systems.py b/tests/test_systems.py index 30dc4ed5e..9541a2034 100644 --- a/tests/test_systems.py +++ b/tests/test_systems.py @@ -7,7 +7,7 @@ @og_test -def test_system_clear(env): +def test_system_clear(env, pipeline_mode): breakfast_table = env.scene.object_registry("name", "breakfast_table") for system_name, system_class in SYSTEM_EXAMPLES.items(): for _ in range(3): diff --git a/tests/test_transition_rules.py b/tests/test_transition_rules.py index ddece38da..51d80204f 100644 --- a/tests/test_transition_rules.py +++ b/tests/test_transition_rules.py @@ -22,7 +22,7 @@ @og_test -def test_dryer_rule(env): +def test_dryer_rule(env, pipeline_mode): assert len(REGISTERED_RULES) > 0, "No rules registered!" clothes_dryer = env.scene.object_registry("name", "clothes_dryer") remover_dishtowel = env.scene.object_registry("name", "remover_dishtowel") @@ -78,7 +78,7 @@ def test_dryer_rule(env): @og_test -def test_washer_rule(env): +def test_washer_rule(env, pipeline_mode): assert len(REGISTERED_RULES) > 0, "No rules registered!" baking_sheet = env.scene.object_registry("name", "baking_sheet") washer = env.scene.object_registry("name", "washer") @@ -149,7 +149,7 @@ def test_washer_rule(env): @og_test -def test_slicing_rule(env): +def test_slicing_rule(env, pipeline_mode): assert len(REGISTERED_RULES) > 0, "No rules registered!" apple = env.scene.object_registry("name", "apple") table_knife = env.scene.object_registry("name", "table_knife") @@ -202,7 +202,7 @@ def test_slicing_rule(env): @og_test -def test_dicing_rule_cooked(env): +def test_dicing_rule_cooked(env, pipeline_mode): assert len(REGISTERED_RULES) > 0, "No rules registered!" half_apple = env.scene.object_registry("name", "half_apple") table_knife = env.scene.object_registry("name", "table_knife") @@ -254,7 +254,7 @@ def test_dicing_rule_cooked(env): @og_test -def test_dicing_rule_uncooked(env): +def test_dicing_rule_uncooked(env, pipeline_mode): assert len(REGISTERED_RULES) > 0, "No rules registered!" half_apple = env.scene.object_registry("name", "half_apple") table_knife = env.scene.object_registry("name", "table_knife") @@ -304,7 +304,7 @@ def test_dicing_rule_uncooked(env): @og_test -def test_melting_rule(env): +def test_melting_rule(env, pipeline_mode): assert len(REGISTERED_RULES) > 0, "No rules registered!" stove = env.scene.object_registry("name", "stove") stockpot = env.scene.object_registry("name", "stockpot") @@ -354,7 +354,7 @@ def test_melting_rule(env): @og_test -def test_cooking_physical_particle_rule_failure_recipe_systems(env): +def test_cooking_physical_particle_rule_failure_recipe_systems(env, pipeline_mode): assert len(REGISTERED_RULES) > 0, "No rules registered!" stove = env.scene.object_registry("name", "stove") stockpot = env.scene.object_registry("name", "stockpot") @@ -394,7 +394,7 @@ def test_cooking_physical_particle_rule_failure_recipe_systems(env): @og_test -def test_cooking_physical_particle_rule_success(env): +def test_cooking_physical_particle_rule_success(env, pipeline_mode): assert len(REGISTERED_RULES) > 0, "No rules registered!" stove = env.scene.object_registry("name", "stove") stockpot = env.scene.object_registry("name", "stockpot") @@ -441,7 +441,7 @@ def test_cooking_physical_particle_rule_success(env): @og_test -def test_mixing_rule_failure_recipe_systems(env): +def test_mixing_rule_failure_recipe_systems(env, pipeline_mode): assert len(REGISTERED_RULES) > 0, "No rules registered!" bowl = env.scene.object_registry("name", "bowl") tablespoon = env.scene.object_registry("name", "tablespoon") @@ -487,7 +487,7 @@ def test_mixing_rule_failure_recipe_systems(env): @og_test -def test_mixing_rule_failure_nonrecipe_systems(env): +def test_mixing_rule_failure_nonrecipe_systems(env, pipeline_mode): assert len(REGISTERED_RULES) > 0, "No rules registered!" bowl = env.scene.object_registry("name", "bowl") tablespoon = env.scene.object_registry("name", "tablespoon") @@ -538,7 +538,7 @@ def test_mixing_rule_failure_nonrecipe_systems(env): @og_test -def test_mixing_rule_success(env): +def test_mixing_rule_success(env, pipeline_mode): assert len(REGISTERED_RULES) > 0, "No rules registered!" bowl = env.scene.object_registry("name", "bowl") tablespoon = env.scene.object_registry("name", "tablespoon") @@ -581,7 +581,7 @@ def test_mixing_rule_success(env): @og_test -def test_cooking_system_rule_failure_recipe_systems(env): +def test_cooking_system_rule_failure_recipe_systems(env, pipeline_mode): assert len(REGISTERED_RULES) > 0, "No rules registered!" stove = env.scene.object_registry("name", "stove") stockpot = env.scene.object_registry("name", "stockpot") @@ -636,7 +636,7 @@ def test_cooking_system_rule_failure_recipe_systems(env): @og_test -def test_cooking_system_rule_failure_nonrecipe_systems(env): +def test_cooking_system_rule_failure_nonrecipe_systems(env, pipeline_mode): assert len(REGISTERED_RULES) > 0, "No rules registered!" stove = env.scene.object_registry("name", "stove") stockpot = env.scene.object_registry("name", "stockpot") @@ -695,7 +695,7 @@ def test_cooking_system_rule_failure_nonrecipe_systems(env): @og_test -def test_cooking_system_rule_failure_nonrecipe_objects(env): +def test_cooking_system_rule_failure_nonrecipe_objects(env, pipeline_mode): assert len(REGISTERED_RULES) > 0, "No rules registered!" stove = env.scene.object_registry("name", "stove") stockpot = env.scene.object_registry("name", "stockpot") @@ -754,7 +754,7 @@ def test_cooking_system_rule_failure_nonrecipe_objects(env): @og_test -def test_cooking_system_rule_success(env): +def test_cooking_system_rule_success(env, pipeline_mode): assert len(REGISTERED_RULES) > 0, "No rules registered!" stove = env.scene.object_registry("name", "stove") stockpot = env.scene.object_registry("name", "stockpot") @@ -818,7 +818,7 @@ def test_cooking_system_rule_success(env): @og_test -def test_cooking_object_rule_failure_wrong_container(env): +def test_cooking_object_rule_failure_wrong_container(env, pipeline_mode): assert len(REGISTERED_RULES) > 0, "No rules registered!" oven = env.scene.object_registry("name", "oven") @@ -861,7 +861,7 @@ def test_cooking_object_rule_failure_wrong_container(env): @og_test -def test_cooking_object_rule_failure_recipe_objects(env): +def test_cooking_object_rule_failure_recipe_objects(env, pipeline_mode): assert len(REGISTERED_RULES) > 0, "No rules registered!" oven = env.scene.object_registry("name", "oven") @@ -903,7 +903,7 @@ def test_cooking_object_rule_failure_recipe_objects(env): @og_test -def test_cooking_object_rule_failure_unary_states(env): +def test_cooking_object_rule_failure_unary_states(env, pipeline_mode): assert len(REGISTERED_RULES) > 0, "No rules registered!" oven = env.scene.object_registry("name", "oven") @@ -946,7 +946,7 @@ def test_cooking_object_rule_failure_unary_states(env): @og_test -def test_cooking_object_rule_failure_binary_system_states(env): +def test_cooking_object_rule_failure_binary_system_states(env, pipeline_mode): assert len(REGISTERED_RULES) > 0, "No rules registered!" oven = env.scene.object_registry("name", "oven") @@ -989,7 +989,7 @@ def test_cooking_object_rule_failure_binary_system_states(env): @og_test -def test_cooking_object_rule_failure_binary_object_states(env): +def test_cooking_object_rule_failure_binary_object_states(env, pipeline_mode): assert len(REGISTERED_RULES) > 0, "No rules registered!" oven = env.scene.object_registry("name", "oven") @@ -1032,7 +1032,7 @@ def test_cooking_object_rule_failure_binary_object_states(env): @og_test -def test_cooking_object_rule_failure_wrong_heat_source(env): +def test_cooking_object_rule_failure_wrong_heat_source(env, pipeline_mode): assert len(REGISTERED_RULES) > 0, "No rules registered!" stove = env.scene.object_registry("name", "stove") @@ -1078,7 +1078,7 @@ def test_cooking_object_rule_failure_wrong_heat_source(env): @og_test -def test_cooking_object_rule_success(env): +def test_cooking_object_rule_success(env, pipeline_mode): assert len(REGISTERED_RULES) > 0, "No rules registered!" oven = env.scene.object_registry("name", "oven") @@ -1148,7 +1148,7 @@ def test_cooking_object_rule_success(env): @og_test -def test_single_toggleable_machine_rule_output_system_failure_wrong_container(env): +def test_single_toggleable_machine_rule_output_system_failure_wrong_container(env, pipeline_mode): assert len(REGISTERED_RULES) > 0, "No rules registered!" food_processor = env.scene.object_registry("name", "food_processor") ice_cream = env.scene.object_registry("name", "scoop_of_ice_cream") @@ -1198,7 +1198,7 @@ def test_single_toggleable_machine_rule_output_system_failure_wrong_container(en @og_test -def test_single_toggleable_machine_rule_output_system_failure_recipe_systems(env): +def test_single_toggleable_machine_rule_output_system_failure_recipe_systems(env, pipeline_mode): assert len(REGISTERED_RULES) > 0, "No rules registered!" blender = env.scene.object_registry("name", "blender") ice_cream = env.scene.object_registry("name", "scoop_of_ice_cream") @@ -1246,7 +1246,7 @@ def test_single_toggleable_machine_rule_output_system_failure_recipe_systems(env @og_test -def test_single_toggleable_machine_rule_output_system_failure_recipe_objects(env): +def test_single_toggleable_machine_rule_output_system_failure_recipe_objects(env, pipeline_mode): assert len(REGISTERED_RULES) > 0, "No rules registered!" blender = env.scene.object_registry("name", "blender") ice_cream = env.scene.object_registry("name", "scoop_of_ice_cream") @@ -1286,7 +1286,7 @@ def test_single_toggleable_machine_rule_output_system_failure_recipe_objects(env @og_test -def test_single_toggleable_machine_rule_output_system_failure_nonrecipe_systems(env): +def test_single_toggleable_machine_rule_output_system_failure_nonrecipe_systems(env, pipeline_mode): assert len(REGISTERED_RULES) > 0, "No rules registered!" blender = env.scene.object_registry("name", "blender") ice_cream = env.scene.object_registry("name", "scoop_of_ice_cream") @@ -1337,7 +1337,7 @@ def test_single_toggleable_machine_rule_output_system_failure_nonrecipe_systems( @og_test -def test_single_toggleable_machine_rule_output_system_failure_nonrecipe_objects(env): +def test_single_toggleable_machine_rule_output_system_failure_nonrecipe_objects(env, pipeline_mode): assert len(REGISTERED_RULES) > 0, "No rules registered!" blender = env.scene.object_registry("name", "blender") ice_cream = env.scene.object_registry("name", "scoop_of_ice_cream") @@ -1387,7 +1387,7 @@ def test_single_toggleable_machine_rule_output_system_failure_nonrecipe_objects( @og_test -def test_single_toggleable_machine_rule_output_system_success(env): +def test_single_toggleable_machine_rule_output_system_success(env, pipeline_mode): assert len(REGISTERED_RULES) > 0, "No rules registered!" blender = env.scene.object_registry("name", "blender") ice_cream = env.scene.object_registry("name", "scoop_of_ice_cream") @@ -1436,7 +1436,7 @@ def test_single_toggleable_machine_rule_output_system_success(env): @og_test -def test_single_toggleable_machine_rule_output_object_failure_unary_states(env): +def test_single_toggleable_machine_rule_output_object_failure_unary_states(env, pipeline_mode): assert len(REGISTERED_RULES) > 0, "No rules registered!" electric_mixer = env.scene.object_registry("name", "electric_mixer") raw_egg = env.scene.object_registry("name", "raw_egg") @@ -1510,7 +1510,7 @@ def test_single_toggleable_machine_rule_output_object_failure_unary_states(env): @og_test -def test_single_toggleable_machine_rule_output_object_success(env): +def test_single_toggleable_machine_rule_output_object_success(env, pipeline_mode): assert len(REGISTERED_RULES) > 0, "No rules registered!" electric_mixer = env.scene.object_registry("name", "electric_mixer") raw_egg = env.scene.object_registry("name", "raw_egg") diff --git a/tests/utils.py b/tests/utils.py index 0d4aeb0a4..085df88ed 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -1,5 +1,6 @@ import math +import pytest import torch as th import omnigibson as og @@ -22,10 +23,11 @@ def og_test(func): - def wrapper(): - assert_test_env() + @pytest.mark.parametrize("pipeline_mode", ["cpu", "cuda"], indirect=True) + def wrapper(pipeline_mode): + assert_test_env(pipeline=pipeline_mode) try: - func(env) + func(env, pipeline_mode) finally: env.scene.reset() @@ -68,26 +70,48 @@ def get_obj_cfg( } -def assert_test_env(): +def assert_test_env(pipeline): global env if env is None: cfg = { + "env": { + "device": pipeline, + }, "scene": { "type": "Scene", }, "objects": [ get_obj_cfg("breakfast_table", "breakfast_table", "skczfi"), get_obj_cfg("bottom_cabinet", "bottom_cabinet", "immwzb"), - get_obj_cfg("dishtowel", "dishtowel", "dtfspn", prim_type=PrimType.CLOTH, abilities={"cloth": {}}), - get_obj_cfg("carpet", "carpet", "ctclvd", prim_type=PrimType.CLOTH, abilities={"cloth": {}}), + ( + get_obj_cfg("dishtowel", "dishtowel", "dtfspn") + if pipeline == "cpu" + else get_obj_cfg( + "dishtowel", "dishtowel", "dtfspn", prim_type=PrimType.CLOTH, abilities={"cloth": {}} + ) + ), + ( + get_obj_cfg("carpet", "carpet", "ctclvd") + if pipeline == "cpu" + else get_obj_cfg("carpet", "carpet", "ctclvd", prim_type=PrimType.CLOTH, abilities={"cloth": {}}) + ), get_obj_cfg("bowl", "bowl", "ajzltc"), get_obj_cfg("bagel", "bagel", "zlxkry", abilities=TEMP_RELATED_ABILITIES), - get_obj_cfg( - "cookable_dishtowel", - "dishtowel", - "dtfspn", - prim_type=PrimType.CLOTH, - abilities={**TEMP_RELATED_ABILITIES, **{"cloth": {}}}, + ( + get_obj_cfg( + "cookable_dishtowel", + "dishtowel", + "dtfspn", + abilities={**TEMP_RELATED_ABILITIES}, + ) + if pipeline == "cpu" + else get_obj_cfg( + "cookable_dishtowel", + "dishtowel", + "dtfspn", + prim_type=PrimType.CLOTH, + abilities={**TEMP_RELATED_ABILITIES, **{"cloth": {}}}, + ) ), get_obj_cfg("microwave", "microwave", "hjjxmi"), get_obj_cfg("stove", "stove", "yhjzwg"), @@ -180,10 +204,7 @@ def assert_test_env(): } if og.sim is None: - # Make sure GPU dynamics are enabled (GPU dynamics needed for cloth) and no flatcache gm.ENABLE_OBJECT_STATES = True - gm.USE_GPU_DYNAMICS = True - gm.ENABLE_FLATCACHE = False gm.ENABLE_TRANSITION_RULES = True else: # Make sure sim is stopped