diff --git a/omnigibson/controllers/controller_base.py b/omnigibson/controllers/controller_base.py index 101b8cec5..b90a56dba 100644 --- a/omnigibson/controllers/controller_base.py +++ b/omnigibson/controllers/controller_base.py @@ -302,8 +302,9 @@ def _dump_state(self): ) def _load_state(self, state): + # Make sure every entry in goal is a numpy array # Load goal - self._goal = state["goal"] + self._goal = None if state["goal"] is None else {name: np.array(goal_state) for name, goal_state in state["goal"].items()} def _serialize(self, state): # Make sure size of the state is consistent, even if we have no goal diff --git a/omnigibson/object_states/adjacency.py b/omnigibson/object_states/adjacency.py index a3cec6f1f..4718b315e 100644 --- a/omnigibson/object_states/adjacency.py +++ b/omnigibson/object_states/adjacency.py @@ -6,7 +6,8 @@ from omnigibson.macros import create_module_macros from omnigibson.object_states.aabb import AABB from omnigibson.object_states.object_state_base import AbsoluteObjectState -from omnigibson.utils.sampling_utils import raytest_batch +from omnigibson.utils.sampling_utils import raytest_batch, raytest +from omnigibson.utils.constants import PrimType # Create settings for this module @@ -60,15 +61,20 @@ def get_equidistant_coordinate_planes(n_planes): return np.stack([first_axes[:, None, :], second_axes[:, None, :]], axis=1) -def compute_adjacencies(obj, axes, max_distance): +def compute_adjacencies(obj, axes, max_distance, use_aabb_center=True): """ Given an object and a list of axes, find the adjacent objects in the axes' positive and negative directions. + If @obj is of PrimType.CLOTH, then adjacent objects are found with respect to the + @obj's centroid particle position + Args: obj (StatefulObject): The object to check adjacencies of. axes (2D-array): (n_axes, 3) array defining the axes to check in. Note that each axis will be checked in both its positive and negative direction. + use_aabb_center (bool): If True and @obj is not of PrimType.CLOTH, will shoot rays from @obj's aabb center. + Otherwise, will dynamically compute starting points based on the requested @axes Returns: list of AxisAdjacencyList: List of length len(axes) containing the adjacencies. @@ -80,17 +86,49 @@ def compute_adjacencies(obj, axes, max_distance): directions[1::2] = -axes # Prepare this object's info for ray casting. - # Use AABB center instead of position because we cannot get valid position - # for fixed objects if fixed links are merged. - aabb_lower, aabb_higher = obj.states[AABB].get_value() - object_position = (aabb_lower + aabb_higher) / 2.0 - prim_paths = obj.link_prim_paths + if obj.prim_type == PrimType.CLOTH: + ray_starts = np.tile(obj.root_link.centroid_particle_position, (len(directions), 1)) + + else: + aabb_lower, aabb_higher = obj.states[AABB].get_value() + object_position = (aabb_lower + aabb_higher) / 2.0 + ray_starts = np.tile(object_position, (len(directions), 1)) + + if not use_aabb_center: + # Dynamically compute start points by iterating over the directions and pre-shooting rays from + # which to shoot back from + # For a given direction, we go in the negative (opposite) direction to the edge of the object extent, + # and then proceed with an additional offset before shooting rays + shooting_offset = 0.01 + + direction_half_extent = directions * (aabb_higher - aabb_lower).reshape(1, 3) / 2.0 + pre_start = object_position.reshape(1, 3) + (direction_half_extent + directions * shooting_offset) + pre_end = object_position.reshape(1, 3) - direction_half_extent + + idx = 0 + obj_link_paths = {link.prim_path for link in obj.links.values()} + def _ray_callback(hit): + # Check for self-hit -- if so, record the position and terminate early + should_continue = True + if hit.rigid_body in obj_link_paths: + ray_starts[idx] = np.array(hit.position) + should_continue = False + return should_continue + + for ray_start, ray_end in zip(pre_start, pre_end): + raytest( + start_point=ray_start, + end_point=ray_end, + only_closest=False, + callback=_ray_callback, + ) + idx += 1 # Prepare the rays to cast. - ray_starts = np.tile(object_position, (len(directions), 1)) ray_endpoints = ray_starts + (directions * max_distance) # Cast time. + prim_paths = obj.link_prim_paths ray_results = raytest_batch( ray_starts, ray_endpoints, @@ -130,7 +168,7 @@ class VerticalAdjacency(AbsoluteObjectState): def _get_value(self): # Call the adjacency computation with th Z axis. - bodies_by_axis = compute_adjacencies(self.obj, np.array([[0, 0, 1]]), m.MAX_DISTANCE_VERTICAL) + bodies_by_axis = compute_adjacencies(self.obj, np.array([[0, 0, 1]]), m.MAX_DISTANCE_VERTICAL, use_aabb_center=False) # Return the adjacencies from the only axis we passed in. return bodies_by_axis[0] @@ -168,7 +206,7 @@ def _get_value(self): coordinate_planes = get_equidistant_coordinate_planes(m.HORIZONTAL_AXIS_COUNT) # Flatten the axis dimension and input into compute_adjacencies. - bodies_by_axis = compute_adjacencies(self.obj, coordinate_planes.reshape(-1, 3), m.MAX_DISTANCE_HORIZONTAL) + bodies_by_axis = compute_adjacencies(self.obj, coordinate_planes.reshape(-1, 3), m.MAX_DISTANCE_HORIZONTAL, use_aabb_center=True) # Now reshape the bodies_by_axis to group by coordinate planes. bodies_by_plane = list(zip(bodies_by_axis[::2], bodies_by_axis[1::2])) diff --git a/omnigibson/prims/cloth_prim.py b/omnigibson/prims/cloth_prim.py index 20927fddb..4e8c22bac 100644 --- a/omnigibson/prims/cloth_prim.py +++ b/omnigibson/prims/cloth_prim.py @@ -61,6 +61,7 @@ def __init__( load_config=None, ): # Internal vars stored + self._centroid_idx = None self._keypoint_idx = None self._keyface_idx = None @@ -115,6 +116,12 @@ def _post_load(self): break assert success, f"Did not adequately subsample keypoints for cloth {self.name}!" + # Compute centroid particle idx based on AABB + aabb_min, aabb_max = np.min(positions, axis=0), np.max(positions, axis=0) + aabb_center = (aabb_min + aabb_max) / 2.0 + dists = np.linalg.norm(positions - aabb_center.reshape(1, 3), axis=-1) + self._centroid_idx = np.argmin(dists) + def _initialize(self): super()._initialize() # TODO (eric): hacky way to get cloth rendering to work (otherwise, there exist some rendering artifacts). @@ -243,6 +250,16 @@ def keypoint_particle_positions(self): """ return self.compute_particle_positions(idxs=self._keypoint_idx) + @property + def centroid_particle_position(self): + """ + Grabs the individual particle that was pre-computed to be the closest to the centroid of this cloth prim. + + Returns: + np.array: centroid particle's (x,y,z) cartesian coordinates relative to the world frame + """ + return self.compute_particle_positions(idxs=[self._centroid_idx])[0] + @property def particle_velocities(self): """ @@ -534,12 +551,10 @@ def _load_state(self, state): # Set values appropriately self._n_particles = state["n_particles"] - for attr in ("positions", "velocities"): - attr_name = f"particle_{attr}" - # Make sure the loaded state is a numpy array, it could have been accidentally casted into a list during - # JSON-serialization - attr_val = np.array(state[attr_name]) if not isinstance(attr_name, np.ndarray) else state[attr_name] - setattr(self, attr_name, attr_val) + # Make sure the loaded state is a numpy array, it could have been accidentally casted into a list during + # JSON-serialization + self.particle_velocities = np.array(state["particle_velocities"]) if not isinstance(state["particle_velocities"], np.ndarray) else state["particle_velocities"] + self.set_particle_positions(positions=np.array(state["particle_positions"]) if not isinstance(state["particle_positions"], np.ndarray) else state["particle_positions"]) def _serialize(self, state): # Run super first diff --git a/omnigibson/prims/joint_prim.py b/omnigibson/prims/joint_prim.py index cd4fba84e..19d64bb8d 100644 --- a/omnigibson/prims/joint_prim.py +++ b/omnigibson/prims/joint_prim.py @@ -422,7 +422,8 @@ def friction(self): Returns: float: friction for this joint """ - return self._articulation_view.get_friction_coefficients(joint_indices=self.dof_indices)[0][0] + return self._articulation_view.get_friction_coefficients(joint_indices=self.dof_indices)[0][0] \ + if og.sim.is_playing() else self.get_attribute("physxJoint:jointFriction") @friction.setter def friction(self, friction): @@ -432,7 +433,9 @@ def friction(self, friction): Args: friction (float): friction to set """ - self._articulation_view.set_friction_coefficients(np.array([[friction]]), joint_indices=self.dof_indices) + self.set_attribute("physxJoint:jointFriction", friction) + if og.sim.is_playing(): + self._articulation_view.set_friction_coefficients(np.array([[friction]]), joint_indices=self.dof_indices) @property def lower_limit(self): diff --git a/omnigibson/utils/sampling_utils.py b/omnigibson/utils/sampling_utils.py index a5377faac..21646cf79 100644 --- a/omnigibson/utils/sampling_utils.py +++ b/omnigibson/utils/sampling_utils.py @@ -251,7 +251,7 @@ def sample_origin_positions(mins, maxes, count, bimodal_mean_fraction, bimodal_s return results -def raytest_batch(start_points, end_points, only_closest=True, ignore_bodies=None, ignore_collisions=None): +def raytest_batch(start_points, end_points, only_closest=True, ignore_bodies=None, ignore_collisions=None, callback=None): """ Computes raytest collisions for a set of rays cast from @start_points to @end_points. @@ -265,6 +265,9 @@ def raytest_batch(start_points, end_points, only_closest=True, ignore_bodies=Non whose collisions should be ignored ignore_collisions (None or list of str): If specified, specifies absolute USD paths to collision geoms whose collisions should be ignored + callback (None or function): If specified and @only_closest is False, the custom callback to use per-hit. + This can be efficient if raytests are meant to terminate early. If None, no custom callback will be used. + Expected signature is callback(hit) -> bool, which returns True if the raycast should continue or not Returns: list of dict or list of list of dict: Results for all rays, where each entry corresponds to the result for the @@ -289,6 +292,7 @@ def raytest_batch(start_points, end_points, only_closest=True, ignore_bodies=Non only_closest=only_closest, ignore_bodies=ignore_bodies, ignore_collisions=ignore_collisions, + callback=callback, )) return results @@ -300,6 +304,7 @@ def raytest( only_closest=True, ignore_bodies=None, ignore_collisions=None, + callback=None, ): """ Computes raytest collision for ray cast from @start_point to @end_point @@ -312,6 +317,9 @@ def raytest( whose collisions should be ignored ignore_collisions (None or list of str): If specified, specifies absolute USD paths to collision geoms whose collisions should be ignored + callback (None or function): If specified and @only_closest is False, the custom callback to use per-hit. + This can be efficient if raytests are meant to terminate early. If None, no custom callback will be used. + Expected signature is callback(hit) -> bool, which returns True if the raycast should continue or not Returns: dict or list of dict: Results for this raytest. If @only_closest=True, then we only return the information from @@ -346,7 +354,7 @@ def raytest( ignore_bodies = set() if ignore_bodies is None else set(ignore_bodies) ignore_collisions = set() if ignore_collisions is None else set(ignore_collisions) - def callback(hit): + def hit_callback(hit): # Only add to hits if we're not ignoring this body or collision if hit.rigid_body not in ignore_bodies and hit.collision not in ignore_collisions: hits.append({ @@ -358,14 +366,14 @@ def callback(hit): "rigidBody": hit.rigid_body, }) # We always want to continue traversing to collect all hits - return True + return True if callback is None else callback(hit) # Grab all collisions og.sim.psqi.raycast_all( origin=start_point, dir=direction, distance=distance, - reportFn=callback, + reportFn=hit_callback, ) # If we only want the closest, we need to sort these hits, otherwise we return them all diff --git a/tests/test_transition_rules.py b/tests/test_transition_rules.py index dadc563c4..54c206250 100644 --- a/tests/test_transition_rules.py +++ b/tests/test_transition_rules.py @@ -1509,5 +1509,3 @@ def test_single_toggleable_machine_rule_output_object_success(): obj = DatasetObject(**obj_cfg) og.sim.import_object(obj) og.sim.step() - -test_single_toggleable_machine_rule_output_system_failure_wrong_container() \ No newline at end of file