Merge pull request #622 from StanfordVL/fix/adjacency_and_friction

Fix/adjacency and friction

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

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]))

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

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):

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

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()