def load_localized_model(km, model_path, reference_frame):
    if model_path != 'nobilia':
        urdf_model = load_urdf_file(model_path)

        return load_model(km, 
                          model_path, 
                          reference_frame, 
                          *[gm.Position(f'{urdf_model.name}_location_{x}') for x in 'xyz'], 
                          gm.Position(f'{urdf_model.name}_location_yaw'))
    else:
        return load_model(km,
                          model_path, 
                          reference_frame, 
                          *[gm.Position(f'nobilia_location_{x}') for x in 'xyz'], 
                          gm.Position('nobilia_location_yaw'))
    def cb_robot_js(self, js_msg):
        if self._state_cb is None:
            return

        state_update = {}
        for x, name in enumerate(js_msg.name):
            if len(js_msg.position) > x:
                state_update[gm.Position(self.robot_prefix +
                                         (name, ))] = js_msg.position[x]
            if len(js_msg.velocity) > x:
                state_update[gm.Velocity(self.robot_prefix +
                                         (name, ))] = js_msg.velocity[x]

        if self.base_aliases is not None:
            try:
                trans = self.tf_buffer.lookup_transform(
                    self.reference_frame, 'base_footprint', rospy.Time(0))
                state_update[self.base_aliases.ap] = math.acos(
                    trans.transform.rotation.w) * 2
                state_update[
                    self.base_aliases.xp] = trans.transform.translation.x
                state_update[
                    self.base_aliases.yp] = trans.transform.translation.y
                state_update[self.base_aliases.av] = 0
                state_update[self.base_aliases.xv] = 0
                state_update[self.base_aliases.yv] = 0
            except (tf2_ros.LookupException, tf2_ros.ConnectivityException,
                    tf2_ros.ExtrapolationException) as e:
                print(
                    f'Exception raised while looking up {self.reference_frame} -> base_footprint:\n{e}'
                )
                return

        self._state_cb(state_update)
def pr2_setup(km, pr2_path):
    start_pose = {'l_elbow_flex_joint' : -2.1213,
                  'l_shoulder_lift_joint': 1.2963,
                  'l_wrist_flex_joint' : -1.16,
                  'l_shoulder_pan_joint': 0.7,
                  'r_shoulder_pan_joint': -1.0,
                  'r_shoulder_lift_joint': 0.9,
                  'r_upper_arm_roll_joint': -1.2,
                  'r_elbow_flex_joint' : -2.1213,
                  'r_wrist_flex_joint' : -1.05,
                  'r_forearm_roll_joint': 3.14,
                  'r_wrist_roll_joint': 0,
                  'torso_lift_joint'   : 0.16825}
    start_pose = {gm.Position(Path(f'{pr2_path}/{n}')): v for n, v in start_pose.items()}
    joint_symbols = [j.position for j in km.get_data(f'{pr2_path}/joints').values() 
                                if hasattr(j, 'position') and gm.is_symbol(j.position)]
    robot_controlled_symbols = {gm.DiffSymbol(j) for j in joint_symbols} # if 'torso' not in str(j)}
    blacklist = {gm.Velocity(Path(f'{pr2_path}/torso_lift_joint'))}
    return joint_symbols, \
           robot_controlled_symbols, \
           start_pose, \
           km.get_data(pr2_path).links['r_gripper_tool_frame'], \
           blacklist
    resting_pose = {
        # 'l_elbow_flex_joint' : -2.1213,
        # 'l_shoulder_lift_joint': 1.2963,
        # 'l_wrist_flex_joint' : -1.16,
        'r_shoulder_pan_joint': -1.0,
        'r_shoulder_lift_joint': 0.9,
        'r_upper_arm_roll_joint': -1.2,
        'r_elbow_flex_joint': -2.1213,
        'r_wrist_flex_joint': -1.05,
        'r_forearm_roll_joint': 3.14,
        'r_wrist_roll_joint': 0,
        #'torso_lift_joint'   : 0.16825
    }
    resting_pose = {
        gm.Position(Path(f'pr2/{n}')): v
        for n, v in resting_pose.items()
    }

    nav_method = rospy.get_param('~nav_method', 'proj')

    gripper = PR2GripperWrapper('/r_gripper_controller')

    pr2_commander = PR2VelCommandProcessor(Path('pr2'),
                                           '/pr2_vel_controller/command',
                                           robot_controlled_symbols,
                                           '/base_controller/command',
                                           base_symbols,
                                           reference_frame=reference_frame)

    behavior = ROSPushingBehavior(
            robot_controlled_symbols |= {base_joint.l_wheel_vel, base_joint.r_wheel_vel}
            robot_controlled_symbols.difference_update(gm.DiffSymbol(s) for s in [base_joint.x_pos, base_joint.y_pos, base_joint.a_pos])
            blacklist.update(gm.DiffSymbol(s) for s in [base_joint.x_pos, base_joint.y_pos, base_joint.a_pos])
            integration_rules = {
                      base_joint.x_pos: base_joint.x_pos + sym_dt * (base_joint.r_wheel_vel * gm.cos(base_joint.a_pos) * base_joint.wheel_radius * 0.5 + base_joint.l_wheel_vel * gm.cos(base_joint.a_pos) * base_joint.wheel_radius * 0.5),
                      base_joint.y_pos: base_joint.y_pos + sym_dt * (base_joint.r_wheel_vel * gm.sin(base_joint.a_pos) * base_joint.wheel_radius * 0.5 + base_joint.l_wheel_vel * gm.sin(base_joint.a_pos) * base_joint.wheel_radius * 0.5),
                      base_joint.a_pos: base_joint.a_pos + sym_dt * (base_joint.r_wheel_vel * (base_joint.wheel_radius / base_joint.wheel_distance) + base_joint.l_wheel_vel * (- base_joint.wheel_radius / base_joint.wheel_distance))}
        start_pose = {'wrist_roll_joint'   : 0.0,
                      'shoulder_pan_joint' : 0.3,
                      'elbow_flex_joint'   : 1.72,
                      'forearm_roll_joint' : -1.2,
                      'upperarm_roll_joint': -1.57,
                      'wrist_flex_joint'   : 1.66,
                      'shoulder_lift_joint': 1.4,
                      'torso_lift_joint'   : 0.2}
        start_pose = {gm.Position(Path(f'{robot_path}/{n}')): v for n, v in start_pose.items()}
    else:
        joint_symbols, \
        robot_controlled_symbols, \
        start_pose, \
        eef, \
        blacklist = generic_setup(km, robot_path, rospy.get_param('~eef', 'gripper_link'))

    collision_world = km.get_active_geometry(gm.free_symbols(handle.pose).union(gm.free_symbols(eef.pose)))

    # Init step
    grasp_in_handle = gm.dot(gm.translation3(0.04, 0, 0), gm.rotation3_rpy(math.pi * 0.5, 0, math.pi))
    goal_pose       = gm.dot(handle.pose, grasp_in_handle)
    goal_0_pose     = gm.subs(goal_pose, {s: 0 for s in gm.free_symbols(goal_pose)})

    start_state = {s: 0 for s in gm.free_symbols(collision_world)}
def create_nobilia_shelf(km,
                         prefix,
                         origin_pose=gm.eye(4),
                         parent_path=Path('world')):
    km.apply_operation(
        f'create {prefix}',
        ExecFunction(prefix, MarkedArticulatedObject, str(prefix)))

    shelf_height = 0.72
    shelf_width = 0.6
    shelf_body_depth = 0.35

    wall_width = 0.016

    l_prefix = prefix + ('links', )
    geom_body_wall_l = Box(
        l_prefix + ('body', ),
        gm.translation3(0, 0.5 * (shelf_width - wall_width), 0),
        gm.vector3(shelf_body_depth, wall_width, shelf_height))
    geom_body_wall_r = Box(
        l_prefix + ('body', ),
        gm.translation3(0, -0.5 * (shelf_width - wall_width), 0),
        gm.vector3(shelf_body_depth, wall_width, shelf_height))

    geom_body_ceiling = Box(
        l_prefix + ('body', ),
        gm.translation3(0, 0, 0.5 * (shelf_height - wall_width)),
        gm.vector3(shelf_body_depth, shelf_width - wall_width, wall_width))
    geom_body_floor = Box(
        l_prefix + ('body', ),
        gm.translation3(0, 0, -0.5 * (shelf_height - wall_width)),
        gm.vector3(shelf_body_depth, shelf_width - wall_width, wall_width))

    geom_body_shelf_1 = Box(
        l_prefix + ('body', ),
        gm.translation3(0.02, 0, -0.2 * (shelf_height - wall_width)),
        gm.vector3(shelf_body_depth - 0.04, shelf_width - wall_width,
                   wall_width))

    geom_body_shelf_2 = Box(
        l_prefix + ('body', ),
        gm.translation3(0.02, 0, 0.2 * (shelf_height - wall_width)),
        gm.vector3(shelf_body_depth - 0.04, shelf_width - wall_width,
                   wall_width))

    geom_body_back = Box(
        l_prefix + ('body', ),
        gm.translation3(0.5 * (shelf_body_depth - 0.005), 0, 0),
        gm.vector3(0.005, shelf_width - 2 * wall_width,
                   shelf_height - 2 * wall_width))

    shelf_geom = [
        geom_body_wall_l, geom_body_wall_r, geom_body_ceiling, geom_body_floor,
        geom_body_back, geom_body_shelf_1, geom_body_shelf_2
    ]

    rb_body = RigidBody(parent_path,
                        origin_pose,
                        geometry=dict(enumerate(shelf_geom)),
                        collision=dict(enumerate(shelf_geom)))

    geom_panel_top = Box(l_prefix + ('panel_top', ), gm.eye(4),
                         gm.vector3(0.357, 0.595, wall_width))
    geom_panel_bottom = Box(l_prefix + ('panel_bottom', ), gm.eye(4),
                            gm.vector3(0.357, 0.595, wall_width))

    handle_width = 0.16
    handle_depth = 0.05
    handle_diameter = 0.012

    geom_handle_r = Box(
        l_prefix + ('handle', ),
        gm.translation3(0.5 * handle_depth,
                        0.5 * (handle_width - handle_diameter), 0),
        gm.vector3(handle_depth, handle_diameter, handle_diameter))
    geom_handle_l = Box(
        l_prefix + ('handle', ),
        gm.translation3(0.5 * handle_depth,
                        -0.5 * (handle_width - handle_diameter), 0),
        gm.vector3(handle_depth, handle_diameter, handle_diameter))
    geom_handle_bar = Box(
        l_prefix + ('handle', ),
        gm.translation3(handle_depth - 0.5 * handle_diameter, 0, 0),
        gm.vector3(handle_diameter, handle_width - handle_diameter,
                   handle_diameter))

    handle_geom = [geom_handle_l, geom_handle_r, geom_handle_bar]

    # Sketch of mechanism
    #
    #           T ---- a
    #         ----      \  Z
    #       b ..... V    \
    #       |      ...... d
    #    B  |       ------
    #       c ------
    #                L
    #
    # Diagonal V is virtual
    #
    #
    # Angles:
    #   a -> alpha (given)
    #   b -> gamma_1 + gamma_2 = gamma
    #   c -> don't care
    #   d -> delta_1 + delta_2 = delta
    #

    opening_position = gm.Position(prefix + ('door', ))

    # Calibration results
    #
    # Solution top hinge: cost = 0.03709980624159568 [ 0.08762252 -0.01433833  0.2858676   0.00871125]
    # Solution bottom hinge: cost = 0.025004236048128934 [ 0.1072496  -0.01232362  0.27271013  0.00489996]

    # Added 180 deg rotation due to -x being the forward facing side in this model
    top_hinge_in_body_marker = gm.translation3(0.08762252 - 0.015, 0,
                                               -0.01433833)
    top_panel_marker_in_top_hinge = gm.translation3(0.2858676 - 0.003,
                                                    -wall_width + 0.0025,
                                                    0.00871125 - 0.003)
    front_hinge_in_top_panel_maker = gm.translation3(0.1072496 - 0.02, 0,
                                                     -0.01232362 + 0.007)
    bottom_panel_marker_in_front_hinge = gm.translation3(
        0.27271013, 0, 0.00489996)

    # Top hinge - Data taken from observation
    body_marker_in_body = gm.dot(
        gm.rotation3_axis_angle(gm.vector3(0, 0, 1), math.pi),
        gm.translation3(0.5 * shelf_body_depth - 0.062,
                        -0.5 * shelf_width + 0.078, 0.5 * shelf_height))
    top_panel_marker_in_top_panel = gm.translation3(
        geom_panel_top.scale[0] * 0.5 - 0.062,
        -geom_panel_top.scale[1] * 0.5 + 0.062, geom_panel_top.scale[2] * 0.5)
    bottom_panel_marker_in_bottom_panel = gm.translation3(
        geom_panel_bottom.scale[0] * 0.5 - 0.062,
        -geom_panel_bottom.scale[1] * 0.5 + 0.062,
        geom_panel_bottom.scale[2] * 0.5)

    top_hinge_in_body = gm.dot(body_marker_in_body, top_hinge_in_body_marker)
    top_panel_in_top_hinge = gm.dot(
        top_panel_marker_in_top_hinge,
        gm.inverse_frame(top_panel_marker_in_top_panel))
    front_hinge_in_top_panel = gm.dot(top_panel_marker_in_top_panel,
                                      front_hinge_in_top_panel_maker)
    bottom_panel_in_front_hinge = gm.dot(
        bottom_panel_marker_in_front_hinge,
        gm.inverse_frame(bottom_panel_marker_in_bottom_panel))

    # Point a in body reference frame
    point_a = gm.dot(gm.diag(1, 0, 1, 1), gm.pos_of(top_hinge_in_body))
    point_d = gm.point3(-shelf_body_depth * 0.5 + 0.09, 0,
                        shelf_height * 0.5 - 0.192)
    # point_d  = gm.point3(-shelf_body_depth * 0.5 + gm.Symbol('point_d_x'), 0, shelf_height * 0.5 - gm.Symbol('point_d_z'))
    # Zero alpha along the vertical axis
    vec_a_to_d = gm.dot(point_d - point_a)
    alpha = gm.atan2(vec_a_to_d[0], -vec_a_to_d[2]) + opening_position

    top_panel_in_body = gm.dot(
        top_hinge_in_body,  # Translation hinge to body frame
        gm.rotation3_axis_angle(gm.vector3(0, 1, 0), -opening_position +
                                0.5 * math.pi),  # Hinge around y
        top_panel_in_top_hinge)
    front_hinge_in_body = gm.dot(top_panel_in_body, front_hinge_in_top_panel)

    # Point b in top panel reference frame
    point_b_in_top_hinge = gm.pos_of(
        gm.dot(gm.diag(1, 0, 1, 1), front_hinge_in_top_panel,
               top_panel_in_top_hinge))
    point_b = gm.dot(gm.diag(1, 0, 1, 1), gm.pos_of(front_hinge_in_body))
    # Hinge lift arm in body reference frame
    point_c_in_bottom_panel = gm.dot(
        gm.diag(1, 0, 1, 1),
        bottom_panel_marker_in_bottom_panel,
        gm.point3(-0.094, -0.034, -0.072),
        # gm.point3(-gm.Symbol('point_c_x'), -0.034, -gm.Symbol('point_c_z'))
    )
    point_c_in_front_hinge = gm.dot(
        gm.diag(1, 0, 1, 1),
        gm.dot(bottom_panel_in_front_hinge, point_c_in_bottom_panel))
    length_z = gm.norm(point_a - point_d)

    vec_a_to_b = point_b - point_a
    length_t = gm.norm(vec_a_to_b)
    length_b = gm.norm(point_c_in_front_hinge[:3])
    # length_l = gm.Symbol('length_l') # 0.34
    length_l = 0.372

    vec_b_to_d = point_d - point_b
    length_v = gm.norm(vec_b_to_d)
    gamma_1 = inner_triangle_angle(length_t, length_v, length_z)
    gamma_2 = inner_triangle_angle(length_b, length_v, length_l)

    top_panel_offset_angle = gm.atan2(point_b_in_top_hinge[2],
                                      point_b_in_top_hinge[0])
    bottom_offset_angle = gm.atan2(point_c_in_front_hinge[2],
                                   point_c_in_front_hinge[0])

    gamma = gamma_1 + gamma_2

    rb_panel_top = RigidBody(l_prefix + ('body', ),
                             gm.dot(rb_body.pose, top_panel_in_body),
                             top_panel_in_body,
                             geometry={0: geom_panel_top},
                             collision={0: geom_panel_top})

    # old offset: 0.5 * geom_panel_top.scale[2] + 0.03
    tf_bottom_panel = gm.dot(
        front_hinge_in_top_panel,
        gm.rotation3_axis_angle(
            gm.vector3(0, 1, 0),
            math.pi + bottom_offset_angle - top_panel_offset_angle),
        gm.rotation3_axis_angle(gm.vector3(0, -1, 0), gamma),
        bottom_panel_in_front_hinge)

    rb_panel_bottom = RigidBody(l_prefix + ('panel_top', ),
                                gm.dot(rb_panel_top.pose, tf_bottom_panel),
                                tf_bottom_panel,
                                geometry={0: geom_panel_bottom},
                                collision={0: geom_panel_bottom})

    handle_transform = gm.dot(
        gm.translation3(geom_panel_bottom.scale[0] * 0.5 - 0.08, 0,
                        0.5 * wall_width),
        gm.rotation3_axis_angle(gm.vector3(0, 1, 0), -math.pi * 0.5))
    rb_handle = RigidBody(l_prefix + ('panel_bottom', ),
                          gm.dot(rb_panel_bottom.pose, handle_transform),
                          handle_transform,
                          geometry={x: g
                                    for x, g in enumerate(handle_geom)},
                          collision={x: g
                                     for x, g in enumerate(handle_geom)})
    # Only debugging
    point_c = gm.dot(rb_panel_bottom.pose, point_c_in_bottom_panel)
    vec_b_to_c = point_c - point_b

    km.apply_operation(f'create {prefix}/links/body',
                       CreateValue(rb_panel_top.parent, rb_body))
    km.apply_operation(f'create {prefix}/links/panel_top',
                       CreateValue(rb_panel_bottom.parent, rb_panel_top))
    km.apply_operation(
        f'create {prefix}/links/panel_bottom',
        CreateValue(l_prefix + ('panel_bottom', ), rb_panel_bottom))
    km.apply_operation(f'create {prefix}/links/handle',
                       CreateValue(l_prefix + ('handle', ), rb_handle))
    km.apply_operation(
        f'create {prefix}/joints/hinge',
        ExecFunction(
            prefix + Path('joints/hinge'), RevoluteJoint,
            CPath(rb_panel_top.parent), CPath(rb_panel_bottom.parent),
            opening_position, gm.vector3(0, 1, 0), gm.eye(4), 0, 1.84, **{
                f'{opening_position}':
                Constraint(0 - opening_position, 1.84 - opening_position,
                           opening_position),
                f'{gm.DiffSymbol(opening_position)}':
                Constraint(-0.25, 0.25, gm.DiffSymbol(opening_position))
            }))
    m_prefix = prefix + ('markers', )
    km.apply_operation(
        f'create {prefix}/markers/body',
        ExecFunction(m_prefix + ('body', ), Frame,
                     CPath(l_prefix + ('body', )),
                     gm.dot(rb_body.pose,
                            body_marker_in_body), body_marker_in_body))
    km.apply_operation(
        f'create {prefix}/markers/top_panel',
        ExecFunction(m_prefix + ('top_panel', ), Frame,
                     CPath(l_prefix + ('panel_top', )),
                     gm.dot(rb_panel_top.pose, top_panel_marker_in_top_panel),
                     top_panel_marker_in_top_panel))
    km.apply_operation(
        f'create {prefix}/markers/bottom_panel',
        ExecFunction(
            m_prefix + ('bottom_panel', ), Frame,
            CPath(l_prefix + ('panel_bottom', )),
            gm.dot(rb_panel_bottom.pose, bottom_panel_marker_in_bottom_panel),
            bottom_panel_marker_in_bottom_panel))

    return NobiliaDebug(
        [
            top_hinge_in_body,
            gm.dot(
                top_hinge_in_body,
                gm.rotation3_axis_angle(gm.vector3(0, 1, 0),
                                        -opening_position + 0.5 * math.pi),
                top_panel_in_top_hinge, front_hinge_in_top_panel),
            body_marker_in_body,
            gm.dot(rb_panel_top.pose, top_panel_marker_in_top_panel),
            gm.dot(rb_panel_bottom.pose, bottom_panel_marker_in_bottom_panel)
        ], [(point_a, vec_a_to_d), (point_a, vec_a_to_b),
            (point_b, vec_b_to_d), (point_b, vec_b_to_c)],
        {
            'gamma_1':
            gamma_1,
            'gamma_1 check_dot':
            gamma_1 - gm.acos(
                gm.dot_product(-vec_a_to_b / gm.norm(vec_a_to_b),
                               vec_b_to_d / gm.norm(vec_b_to_d))),
            'gamma_1 check_cos':
            gamma_1 - inner_triangle_angle(
                gm.norm(vec_a_to_b), gm.norm(vec_b_to_d), gm.norm(vec_a_to_d)),
            'gamma_2':
            gamma_2,
            'gamma_2 check_dot':
            gamma_2 - gm.acos(
                gm.dot_product(vec_b_to_c / gm.norm(vec_b_to_c),
                               vec_b_to_d / gm.norm(vec_b_to_d))),
            'length_v':
            length_v,
            'length_b':
            length_b,
            'length_l':
            length_l,
            'position':
            opening_position,
            'alpha':
            alpha,
            'dist c d':
            gm.norm(point_d - point_c)
        }, {
            gm.Symbol('point_c_x'): 0.094,
            gm.Symbol('point_c_z'): 0.072,
            gm.Symbol('point_d_x'): 0.09,
            gm.Symbol('point_d_z'): 0.192,
            gm.Symbol('length_l'): 0.372
        })
        else:
            pushing_controller = PushingController(
                km,
                eef_path,
                part_path,
                robot_controlled_symbols,
                start_state,
                cam_path,
                use_geom_circulation,
                None,
                weights,
                collision_avoidance_paths=active_parts_to_avoid)

        if robot in start_poses:
            start_state.update({
                gm.Position(Path(robot) + (k, )): v
                for k, v in start_poses[robot].items()
            })

        start_state.update({
            c.symbol: 0.0
            for c in pushing_controller.controlled_values.values()
        })
        # start_state.update({s: 1.84 for s in gm.free_symbols(obj.pose)})
        start_state.update({s: 0.4 for s in gm.free_symbols(obj.pose)})

        printed_exprs['relative_vel'] = gm.norm(
            pushing_controller.p_internals.relative_pos)

        integrator = CommandIntegrator(pushing_controller.qpb,
                                       integration_rules,
Exemple #8
0
def create_door(km,
                prefix,
                height,
                width,
                frame_width=0.05,
                to_world_tf=gm.eye(4)):
    km.apply_operation(f'create {prefix}',
                       ExecFunction(prefix, ArticulatedObject, 'door'))

    prefix = prefix + ('links', )

    base_plate_geom = Box(prefix + ('frame', ), gm.translation3(0, 0, 0.015),
                          gm.vector3(0.2, width + 0.2, 0.03))
    frame_pillar_l_geom = Box(
        prefix + ('frame', ),
        gm.translation3(0, 0.5 * (width + frame_width), 0.5 * height + 0.03),
        gm.vector3(frame_width, frame_width, height))
    frame_pillar_r_geom = Box(
        prefix + ('frame', ),
        gm.translation3(0, -0.5 * (width + frame_width), 0.5 * height + 0.03),
        gm.vector3(frame_width, frame_width, height))
    frame_bar_geom = Box(
        prefix + ('frame', ),
        gm.translation3(0, 0, height + 0.5 * frame_width + 0.03),
        gm.vector3(frame_width, width + 2 * frame_width, frame_width))
    frame_rb = RigidBody(Path('world'),
                         to_world_tf,
                         geometry={
                             1: base_plate_geom,
                             2: frame_pillar_l_geom,
                             3: frame_pillar_r_geom,
                             4: frame_bar_geom
                         },
                         collision={
                             1: base_plate_geom,
                             2: frame_pillar_l_geom,
                             3: frame_pillar_r_geom,
                             4: frame_bar_geom
                         })
    door_geom1 = Box(prefix + ('door', ), gm.translation3(0.015, 0, 0),
                     gm.vector3(0.03, width, height))
    door_geom2 = Box(prefix + ('door', ), gm.translation3(-0.005, 0, 0.01),
                     gm.vector3(0.01, width + 0.02, height + 0.01))

    handle_bar_geom = Box(prefix + ('handle', ),
                          gm.translation3(-0.08, 0.06, 0),
                          gm.vector3(0.02, 0.12, 0.02))
    handle_cylinder_geom = Cylinder(
        prefix + ('handle', ),
        gm.dot(gm.translation3(-0.04, 0, 0),
               gm.rotation3_axis_angle(gm.vector3(0, 1, 0), 0.5 * math.pi)),
        0.02, 0.08)

    door_rb = RigidBody(Path(f'{prefix}/frame'),
                        gm.translation3(0.0, 0.5 * -width - 0.01, 0),
                        geometry={
                            1: door_geom1,
                            2: door_geom2
                        },
                        collision={
                            1: door_geom1,
                            2: door_geom2
                        })

    handle_rb = RigidBody(Path(f'{prefix}/door'),
                          gm.eye(4),
                          geometry={
                              1: handle_bar_geom,
                              2: handle_cylinder_geom
                          },
                          collision={
                              1: handle_bar_geom,
                              2: handle_cylinder_geom
                          })

    km.apply_operation(f'create {prefix}/frame',
                       CreateValue(prefix + ('frame', ), frame_rb))
    km.apply_operation(f'create {prefix}/door',
                       CreateValue(prefix + ('door', ), door_rb))
    km.apply_operation(f'create {prefix}/handle',
                       CreateValue(prefix + ('handle', ), handle_rb))

    door_position = gm.Position('door')
    handle_position = gm.Position('handle')

    prefix = prefix[:-1] + ('joints', )
    km.apply_operation(
        f'create {prefix}',
        ExecFunction(
            prefix + ('hinge', ), RevoluteJoint, CPath(door_rb.parent),
            CPath(handle_rb.parent), door_position, gm.vector3(0, 0, -1),
            gm.translation3(0.5 * -frame_width - 0.005, 0.5 * width + 0.01,
                            0.5 * height + 0.03), 0, 0.75 * math.pi, 100, 1,
            0))

    km.apply_operation(
        f'create {prefix}',
        ExecFunction(prefix + ('handle', ), RevoluteJoint,
                     CPath(handle_rb.parent), CPath(f'{prefix}/handle'),
                     handle_position, gm.vector3(-1, 0, 0),
                     gm.translation3(0, -0.5 * width - 0.02 + 0.06,
                                     0), 0, 0.25 * math.pi, 100, 1, 0))

    prefix = prefix[:-1]
    km.apply_operation(
        f'connect {prefix}/links/frame {prefix}/links/door',
        CreateURDFFrameConnection(prefix + ('joints', 'hinge'),
                                  Path(door_rb.parent),
                                  Path(handle_rb.parent)))
    km.apply_operation(
        f'connect {prefix}/links/door {prefix}/links/handle',
        CreateURDFFrameConnection(prefix + ('joints', 'handle'),
                                  Path(handle_rb.parent),
                                  Path(f'{prefix}/links/handle')))
    km.apply_operation(
        f'add lock {door_position}',
        ConditionalDoorHandleConstraints(door_position, handle_position,
                                         math.pi * 0.01, math.pi * 0.15))
    visualizer = ROSBPBVisualizer('~visualization', base_frame='world')

    km = GeometryModel()

    with open(
            res_pkg_path(
                'package://kineverse_experiment_world/urdf/iiwa_wsg_50.urdf'),
            'r') as f:
        iiwa_urdf = urdf_filler(
            URDF.from_xml_string(hacky_urdf_parser_fix(f.read())))
        load_urdf(km, Path('iiwa'), iiwa_urdf)

    km.clean_structure()
    km.dispatch_events()

    door_x, door_y = [gm.Position(f'localization_{x}') for x in 'xy']

    create_door(km,
                Path('door'),
                0.5,
                0.35,
                to_world_tf=gm.translation3(door_x, door_y, 0.1))

    km.clean_structure()
    km.dispatch_events()

    door = km.get_data('door')
    handle = door.links['handle']
    iiwa = km.get_data('iiwa')
    symbols = gm.free_symbols(door.links['handle'].pose).union(
        gm.free_symbols(iiwa.links['link_7'].pose))
import kineverse.gradients.gradient_math as gm
import numpy as np
import matplotlib.pyplot as plt

from kineverse.visualization.plotting import hsv_to_rgb, \
            rgb_to_hex

if __name__ == '__main__':

    a, b = [gm.Position(x) for x in 'ab']

    l = 2
    a_in_w = gm.dot(gm.translation3(0, 0, 2), gm.translation3(0, 0, -a))
    d_in_a = gm.rotation3_axis_angle(gm.vector3(0, 1, 0), gm.acos(a / l))
    d_in_w = gm.dot(a_in_w, d_in_a)
    A = gm.pos_of(a_in_w)
    B = gm.dot(d_in_w, gm.point3(0, 0, l))
    C = gm.dot(d_in_w, gm.point3(0, 0, l * 0.5))
    D = gm.dot(d_in_w, gm.point3(0, 0, l * 0.25))
    E = gm.dot(d_in_w, gm.point3(0, 0, l * 0.75))

    lock_bound = gm.alg_not(
        gm.alg_and(gm.less_than(b, 0.3), gm.less_than(1.99, a)))

    # PLOT OF MOVEMENT

    As = []
    Bs = []
    Cs = []
    Ds = []
    Es = []
    eef_path = Path(f'pr2/links/{eef_link}')
    cam_path = Path('pr2/links/head_mount_link')

    resting_pose = {#'l_elbow_flex_joint' : -2.1213,
                    #'l_shoulder_lift_joint': 1.2963,
                    #'l_wrist_flex_joint' : -1.16,
                    'r_shoulder_pan_joint': -1.0,
                    'r_shoulder_lift_joint': 0.9,
                    'r_upper_arm_roll_joint': -1.2,
                    'r_elbow_flex_joint' : -2.1213,
                    'r_wrist_flex_joint' : -1.05,
                    'r_forearm_roll_joint': 0,
                    'r_wrist_roll_joint': 0,
                    # 'torso_lift_joint'   : 0.16825
                    }
    resting_pose = {gm.Position(Path(f'pr2/{n}')): v for n, v in resting_pose.items()}

    # Object stuff
    grasp_poses = {}
    for p, pose_params in body_paths.items():
        if not km.has_data(p):
            print(f'Link {p} is not part of articulation model')
            exit(1)
        if type(pose_params) is not list and type(pose_params) is not tuple:
            print(f'Grasp pose for {p} is not a list.')
            exit(1)
        if len(pose_params) == 6:
            grasp_poses[Path(p)] = gm.frame3_rpy(*pose_params[-3:], gm.point3(*pose_params[:3]))
        elif len(pose_params) == 7:
            grasp_poses[Path(p)] = gm.frame3_quaternion(*pose_params)
        else: