OMPL constrained planning

main.py

import numpy as np
from airo_drake import SingleArmScene, add_floor, add_manipulator, add_meshcat, finish_build, animate_joint_trajectory, \
    time_parametrize_toppra, visualize_frame, add_wall
from airo_planner import SingleArmOmplPlanner, function_to_ompl, uniform_symmetric_joint_bounds
from airo_typing import HomogeneousMatrixType, JointConfigurationType
from ompl import base as ob
from ompl import geometric as og
from pydrake.planning import RobotDiagramBuilder, SceneGraphCollisionChecker
from ur_analytic_ik_ext import ur3e

DOF = 6
START_JOINTS = np.deg2rad([0, -90, 90, -90, -90, 0])
GOAL_JOINTS = np.array([2.32027433, -1.56661515, 1.58208421, -1.58626538, -1.57079633,
                        -0.82131832])  # np.deg2rad([0, -90, -90, -90, 90, 0])


def bounds_to_ompl(joint_bounds) -> ob.RealVectorBounds:
    degrees_of_freedom = len(joint_bounds[0])
    bounds = ob.RealVectorBounds(degrees_of_freedom)
    joint_bounds_lower = joint_bounds[0]
    joint_bounds_upper = joint_bounds[1]
    for i in range(degrees_of_freedom):
        bounds.setLow(i, joint_bounds_lower[i])
        bounds.setHigh(i, joint_bounds_upper[i])
    return bounds


def _get_space():
    space = ob.RealVectorStateSpace(DOF)
    space.setBounds(bounds_to_ompl(uniform_symmetric_joint_bounds(6)))
    return space


def softplus(x, steepness, threshold):
    return np.log(1 + np.exp(steepness * (-x - threshold)))


class TCPConstraint(ob.Constraint):
    # The TCP should be pointing along the negative z-axis.
    def __init__(self, tolerance: float = 0.05):
        super(TCPConstraint, self).__init__(DOF, 1)

        self.threshold = 1.0 - tolerance

    def function(self, x, out):
        eef_pose = ur3e.forward_kinematics(*x.tolist())
        dotprod = eef_pose[:3, 2].dot(np.array([0., 0., -1.]))
        out[0] = 1.0 - dotprod
        # out[0] = softplus(dotprod, 4, 0.95)
        # eef_pose = ur3e.forward_kinematics(*x.tolist())
        # if eef_pose[:3, 2].dot(np.array([0., 0., -1.])) > self.threshold:
        #     out[0] = 0.
        # else:
        #     out[0] = 1.0 - eef_pose[:3, 2].dot(np.array([0., 0., -1.]))

    # def jacobian(self, x, out):
    #     pass
    # We could compute the Jacobian using Sympy, but that requires modelling forward kinematics as a symbolic function.


def _state_to_ompl(state_numpy: np.ndarray, space: ob.StateSpace) -> ob.State:
    state = ob.State(space)
    for i in range(len(state_numpy)):
        state()[i] = state_numpy[i]
    return state


def _state_from_ompl(state: ob.State, n: int) -> np.ndarray:
    return np.array([state[i] for i in range(n)])


def inverse_kinematics_fn(tcp_pose: HomogeneousMatrixType) -> list[JointConfigurationType]:
    solutions = ur3e.inverse_kinematics(tcp_pose)
    solutions = [solution.squeeze() for solution in solutions]
    return solutions


def constrained():
    robot_diagram_builder = RobotDiagramBuilder()

    meshcat = add_meshcat(robot_diagram_builder)
    arm_index, gripper_index = add_manipulator(robot_diagram_builder, "ur3e", "robotiq_2f_85", static_gripper=True)
    add_floor(robot_diagram_builder)
    add_wall(robot_diagram_builder, "YZ", 0.2, 0.1, 0.05, np.array([0.0, -0.3, 0.1]))
    robot_diagram, context = finish_build(robot_diagram_builder)

    scene = SingleArmScene(robot_diagram, arm_index, gripper_index, meshcat)

    collision_checker = SceneGraphCollisionChecker(
        model=scene.robot_diagram,
        robot_model_instances=[scene.arm_index, scene.gripper_index],
        edge_step_size=0.125,  # Arbitrary value: we don't use the CheckEdgeCollisionFree
        env_collision_padding=0.005,
        self_collision_padding=0.005,
    )

    plant = scene.robot_diagram.plant()
    plant_context = plant.GetMyContextFromRoot(context)
    plant.SetPositions(plant_context, scene.arm_index, START_JOINTS)
    scene.robot_diagram.ForcedPublish(context)

    space = _get_space()
    constraint = TCPConstraint()
    constraint_tolerance = 0.05
    constraint.setTolerance(constraint_tolerance)

    css = ob.ProjectedStateSpace(space, constraint)
    csi = ob.ConstrainedSpaceInformation(css)

    delta = ob.CONSTRAINED_STATE_SPACE_DELTA
    lambda_ = ob.CONSTRAINED_STATE_SPACE_LAMBDA
    css.setup()
    css.setDelta(delta)
    css.setLambda(lambda_)
    simple_setup = og.SimpleSetup(csi)

    planner = og.RRTConnect(simple_setup.getSpaceInformation())
    simple_setup.setPlanner(planner)

    is_state_valid_fn = function_to_ompl(collision_checker.CheckConfigCollisionFree, DOF)
    simple_setup.setStateValidityChecker(ob.StateValidityCheckerFn(is_state_valid_fn))
    simple_setup.setOptimizationObjective(ob.PathLengthOptimizationObjective(simple_setup.getSpaceInformation()))
    step = float(np.deg2rad(5))
    resolution = step / space.getMaximumExtent()
    simple_setup.getSpaceInformation().setStateValidityCheckingResolution(resolution)

    goal_pose = np.array([[0.0, -1.0, 0.0, 0.30],
                          [-1.0, 0.0, 0.0, -0.13],
                          [0.0, 0.0, -1.0, 0.30],
                          [0.0, 0.0, 0.0, 1.0]])
    visualize_frame(scene.meshcat, "goal_pose", goal_pose)
    input("Goal pose frame is now shown.")
    goal_joints_solutions = inverse_kinematics_fn(goal_pose)
    print("Found", len(goal_joints_solutions), "solutions.")

    goal_joints_solutions = [solution for solution in goal_joints_solutions if
                             collision_checker.CheckConfigCollisionFree(solution)]
    print("Found", len(goal_joints_solutions), "solutions that are not in collision.")

    print("Checking constraint function values for goal poses (0 = satisfies constraints).")
    for goal_joints in goal_joints_solutions:
        out = np.empty(1)
        constraint.function(goal_joints, out)
        print(goal_joints, out, "✅" if out[0] < constraint_tolerance else "💥")

    for goal_joints in goal_joints_solutions:
        simple_setup.clear()
        space = simple_setup.getStateSpace()
        start_state = _state_to_ompl(START_JOINTS, space)
        goal_state = _state_to_ompl(goal_joints, space)
        simple_setup.setStartAndGoalStates(start_state, goal_state)

        simple_setup.solve()

        if not simple_setup.haveExactSolutionPath():
            print("No exact solution found.")
            continue

        # Simplify and smooth the path, note that this conserves path validity
        simple_setup.simplifySolution()
        path_simplifier = og.PathSimplifier(simple_setup.getSpaceInformation())
        path = simple_setup.getSolutionPath()
        path_simplifier.smoothBSpline(path)

        # Extract path
        path = np.array([_state_from_ompl(path.getState(i), DOF) for i in range(path.getStateCount())])

        trajectory = time_parametrize_toppra(scene.robot_diagram.plant(), path)

        animate_joint_trajectory(scene.meshcat, scene.robot_diagram, scene.arm_index, trajectory)

        input("Next / Finish?")


def unconstrained():
    robot_diagram_builder = RobotDiagramBuilder()

    meshcat = add_meshcat(robot_diagram_builder)
    arm_index, gripper_index = add_manipulator(robot_diagram_builder, "ur3e", "robotiq_2f_85", static_gripper=True)
    add_floor(robot_diagram_builder)
    add_wall(robot_diagram_builder, "YZ", 0.2, 0.1, 0.05, np.array([0.0, -0.3, 0.1]))
    robot_diagram, context = finish_build(robot_diagram_builder)

    scene = SingleArmScene(robot_diagram, arm_index, gripper_index, meshcat)

    visualize_frame(scene.meshcat, "goal_pose", ur3e.forward_kinematics(*GOAL_JOINTS.tolist()))

    collision_checker = SceneGraphCollisionChecker(
        model=scene.robot_diagram,
        robot_model_instances=[scene.arm_index, scene.gripper_index],
        edge_step_size=0.125,  # Arbitrary value: we don't use the CheckEdgeCollisionFree
        env_collision_padding=0.005,
        self_collision_padding=0.005,
    )

    plant = scene.robot_diagram.plant()
    plant_context = plant.GetMyContextFromRoot(context)
    plant.SetPositions(plant_context, scene.arm_index, START_JOINTS)
    scene.robot_diagram.ForcedPublish(context)

    planner = SingleArmOmplPlanner(collision_checker.CheckConfigCollisionFree,
                                   inverse_kinematics_fn=inverse_kinematics_fn)

    path = planner.plan_to_joint_configuration(START_JOINTS, GOAL_JOINTS)

    for q in path:
        eef_pose = ur3e.forward_kinematics(*q.tolist())
        print(q, eef_pose[:3, 2].dot(np.array([0., 0., -1.])))

    trajectory = time_parametrize_toppra(scene.robot_diagram.plant(), path)

    animate_joint_trajectory(scene.meshcat, scene.robot_diagram, scene.arm_index, trajectory)

    input("Finish?")


def visualize_pose():
    robot_diagram_builder = RobotDiagramBuilder()

    meshcat = add_meshcat(robot_diagram_builder)
    arm_index, gripper_index = add_manipulator(robot_diagram_builder, "ur3e", "robotiq_2f_85", static_gripper=True)
    add_floor(robot_diagram_builder)
    robot_diagram, context = finish_build(robot_diagram_builder)

    scene = SingleArmScene(robot_diagram, arm_index, gripper_index, meshcat)

    # joints = [2.32027433, -0.45024119, -0.23098142, 2.25201893, 1.57079633, 2.32027433]
    # joints = [ 2.32027433, -0.66580611,  0.23098142,  2.00562102,  1.57079633 , 2.32027433]
    joints = [2.32027433, -1.56661515, 1.58208421, -1.58626538, -1.57079633, -0.82131832]
    # [ 0.00350267 -1.5749775  -1.58208421 -1.55532727  1.57079633  0.00350267] [-2.22044605e-16] 💥
    # [ 0.00350267 -2.47578654 -0.23098142  1.13597163 -1.57079633 -3.13808999] [0.] ✅

    plant = scene.robot_diagram.plant()
    plant_context = plant.GetMyContextFromRoot(context)
    plant.SetPositions(plant_context, scene.arm_index, joints)
    scene.robot_diagram.ForcedPublish(context)

    input("Check meshcat")


if __name__ == '__main__':
    # unconstrained()
    constrained()
    # visualize_pose()