Skip to content

object_state_utils

sample_kinematics(predicate, objA, objB, use_ray_casting_method=False, max_trials=10, z_offset=0.05, skip_falling=False)

Samples the given @predicate kinematic state for @objA with respect to @objB

Parameters:

Name Type Description Default
predicate str

Name of the predicate to sample, e.g.: "onTop"

required
objA StatefulObject

Object whose state should be sampled. e.g.: for sampling a microwave on a cabinet, @objA is the microwave

required
objB StatefulObject

Object who is the reference point for @objA's state. e.g.: for sampling a microwave on a cabinet, @objB is the cabinet

required
use_ray_casting_method bool

Whether to use raycasting for sampling or not

False
max_trials int

Number of attempts for sampling

10
z_offset float

Z-offset to apply to the sampled pose

0.05
skip_falling bool

Whether to let @objA fall after its position is sampled or not

False

Returns:

Name Type Description
bool

True if successfully sampled, else False

Source code in utils/object_state_utils.py
def sample_kinematics(
    predicate,
    objA,
    objB,
    use_ray_casting_method=False,
    max_trials=10,
    z_offset=0.05,
    skip_falling=False,
):
    """
    Samples the given @predicate kinematic state for @objA with respect to @objB

    Args:
        predicate (str): Name of the predicate to sample, e.g.: "onTop"
        objA (StatefulObject): Object whose state should be sampled. e.g.: for sampling a microwave
            on a cabinet, @objA is the microwave
        objB (StatefulObject): Object who is the reference point for @objA's state. e.g.: for sampling
            a microwave on a cabinet, @objB is the cabinet
        use_ray_casting_method (bool): Whether to use raycasting for sampling or not
        max_trials (int): Number of attempts for sampling
        z_offset (float): Z-offset to apply to the sampled pose
        skip_falling (bool): Whether to let @objA fall after its position is sampled or not

    Returns:
        bool: True if successfully sampled, else False
    """

    assert z_offset > 0.5 * 9.81 * (og.sim.get_physics_dt() ** 2) + 0.02,\
        f"z_offset {z_offset} is too small for the current physics_dt {og.sim.get_physics_dt()}"

    # Run import here to avoid circular imports
    # No supporting surface annotation found, fallback to use ray-casting
    from omnigibson.objects.dataset_object import DatasetObject
    if (
        not isinstance(objB, DatasetObject) or
        len(objB.supporting_surfaces) == 0 or
        predicate not in objB.supporting_surfaces
    ):
        use_ray_casting_method = True

    # Wake objects accordingly and make sure both are kept still
    objA.wake()
    objB.wake()

    objA.keep_still()
    objB.keep_still()

    # Save the state of the simulator
    state = og.sim.dump_state()

    # Attempt sampling
    for i in range(max_trials):
        pos = None
        if hasattr(objA, "orientations") and objA.orientations is not None:
            orientation = objA.sample_orientation()
        else:
            orientation = np.array([0, 0, 0, 1.0])

        # Orientation needs to be set for stable_z_on_aabb to work correctly
        # Position needs to be set to be very far away because the object's
        # original position might be blocking rays (use_ray_casting_method=True)
        old_pos = np.array([200, 200, 200])
        objA.set_position_orientation(old_pos, orientation)
        objA.keep_still()
        # We also need to step physics to make sure the pose propagates downstream (e.g.: to Bounding Box computations)
        og.sim.step_physics()

        # This would slightly change because of the step_physics call.
        old_pos, orientation = objA.get_position_orientation()

        if use_ray_casting_method:
            if predicate == "onTop":
                params = m.ON_TOP_RAY_CASTING_SAMPLING_PARAMS
            elif predicate == "inside":
                params = m.INSIDE_RAY_CASTING_SAMPLING_PARAMS
            else:
                raise ValueError(f"predicate must be either onTop or inside in order to use ray casting-based "
                                 f"kinematic sampling, but instead got: {predicate}")

            # Run import here to avoid circular imports
            from omnigibson.objects.dataset_object import DatasetObject
            if isinstance(objA, DatasetObject):
                # Retrieve base CoM frame-aligned bounding box parallel to the XY plane
                parallel_bbox_center, parallel_bbox_orn, parallel_bbox_extents, _ = objA.get_base_aligned_bbox(
                    xy_aligned=True
                )
            else:
                aabb_lower, aabb_upper = objA.states[AABB].get_value()
                parallel_bbox_center = (aabb_lower + aabb_upper) / 2.0
                parallel_bbox_orn = np.array([0.0, 0.0, 0.0, 1.0])
                parallel_bbox_extents = aabb_upper - aabb_lower

            sampling_results = sampling_utils.sample_cuboid_on_object_symmetric_bimodal_distribution(
                objB,
                num_samples=1,
                cuboid_dimensions=parallel_bbox_extents,
                axis_probabilities=[0, 0, 1],
                refuse_downwards=True,
                undo_cuboid_bottom_padding=True,
                **params,
            )

            sampled_vector = sampling_results[0][0]
            sampled_quaternion = sampling_results[0][2]

            sampling_success = sampled_vector is not None
            if sampling_success:
                # Move the object from the original parallel bbox to the sampled bbox
                parallel_bbox_rotation = R.from_quat(parallel_bbox_orn)
                sample_rotation = R.from_quat(sampled_quaternion)
                original_rotation = R.from_quat(orientation)

                # The additional orientation to be applied should be the delta orientation
                # between the parallel bbox orientation and the sample orientation
                additional_rotation = sample_rotation * parallel_bbox_rotation.inv()
                combined_rotation = additional_rotation * original_rotation
                orientation = combined_rotation.as_quat()

                # The delta vector between the base CoM frame and the parallel bbox center needs to be rotated
                # by the same additional orientation
                diff = old_pos - parallel_bbox_center
                rotated_diff = additional_rotation.apply(diff)
                pos = sampled_vector + rotated_diff
        else:
            random_idx = np.random.randint(len(objB.supporting_surfaces[predicate].keys()))
            objB_link_name = list(objB.supporting_surfaces[predicate].keys())[random_idx]
            random_height_idx = np.random.randint(len(objB.supporting_surfaces[predicate][objB_link_name]))
            height, height_map = objB.supporting_surfaces[predicate][objB_link_name][random_height_idx]
            obj_half_size = np.max(objA.aabb_extent) / 2 * 100
            obj_half_size_scaled = np.array([obj_half_size / objB.scale[1], obj_half_size / objB.scale[0]])
            obj_half_size_scaled = np.ceil(obj_half_size_scaled).astype(np.int)
            height_map_eroded = cv2.erode(height_map, np.ones(obj_half_size_scaled, np.uint8))

            valid_pos = np.array(height_map_eroded.nonzero())
            if valid_pos.shape[1] != 0:
                random_pos_idx = np.random.randint(valid_pos.shape[1])
                random_pos = valid_pos[:, random_pos_idx]
                y_map, x_map = random_pos
                y = y_map / 100.0 - 2
                x = x_map / 100.0 - 2
                z = height

                pos = np.array([x, y, z])
                pos *= objB.scale

                # the supporting surface is defined w.r.t to the link frame, so we need to convert it into
                # the world frame
                link_pos, link_quat = objB.links[objB_link_name].get_position_orientation()
                pos = T.quat2mat(link_quat).dot(pos) + np.array(link_pos)
                # Get the combined AABB.
                lower, _ = objA.states[AABB].get_value()
                # Move the position to a stable Z for the object.
                pos[2] += objA.get_position()[2] - lower[2]

        if pos is None:
            success = False
        else:
            pos[2] += z_offset
            objA.set_position_orientation(pos, orientation)
            objA.keep_still()

            og.sim.step_physics()
            success = not objA.in_contact()

        if og.debug_sampling:
            print("sample_kinematics", success)
            embed()

        if success:
            break
        else:
            og.sim.load_state(state)

    if success and not skip_falling:
        objA.set_position_orientation(pos, orientation)
        objA.keep_still()

        # Let it fall for 0.2 second
        for _ in range(int(0.2 / og.sim.get_physics_dt())):
            og.sim.step_physics()
            if objA.in_contact():
                break

        # Render at the end
        og.sim.render()

    return success