fruitbot.py

#!/usr/bin/env python3
#
#   fruitbot.py
#
#   Run the fruit ninja robot
#
#   Team Spaghetti
#      John (Jack) Maxfield
#      Angelina Ye
#      Isabella Zhang
#
#   Publish:    /visualization_marker_array     visualization_msgs/MarkerArray
#   Publish:    /joint_states
import rospy

from visualization_msgs.msg import Marker
from visualization_msgs.msg import MarkerArray
import numpy as np
import random
from sensor_msgs.msg   import JointState
from hw6code.kinematics import Kinematics # todo: can't import from this package
                                          # due to hyphen in name

TWO_PI = 2 * np.pi

# finds number equal to q1 mod 2pi
# which is closest to q0
def closest_rdr_mod_2pi(q0, q1):
    q0 /= TWO_PI
    q1 /= TWO_PI
    dif = q1 - q0
    lo = q1 - np.ceil(dif)
    hi = q1 - np.floor(dif)
    # print(f'q0 = {q0}, q1 = {q1}, lo = {lo}, hi = {hi}')
    # assert np.allclose(q1 - lo, np.rint(q1 - lo))
    # assert np.allclose(q1 - hi, np.rint(q1 - hi))
    # assert (q0 - lo) <= 1.001
    # assert (q0 - hi) <= 1.001
    # assert np.allclose(abs(hi - lo), 1)
    return TWO_PI * lo if abs(q0 - lo) < abs(q0 - hi) else TWO_PI * hi

def mod_pi_range(x):
    return ((x + np.pi) % TWO_PI) - np.pi

def interp_away_from_pi(interp_factor, start, end):
    start_0 = mod_pi_range(start)
    end_0 = mod_pi_range(end)
    return (1 - interp_factor) * start_0 + interp_factor * end_0

def interp_joints(q0, q1, i):
    return np.array([
        (1 - i) * q0[0] + i * q1[0],
        (1 - i) * q0[1] + i * q1[1],
        interp_away_from_pi(i, q0[2], q1[2]),
        (1 - i) * q0[3] + i * q1[3],
        (1 - i) * q0[4] + i * q1[4],
        (1 - i) * q0[5] + i * q1[5],
    ])

def totalError(kin, pd, Rd, q):
    p = np.zeros((3,1)) # posn
    R = np.identity(3)  # orientation
    kin.fkin(q, p, R)
    ep = pd - p.reshape(3)
    xd = Rd[:,0]
    yd = Rd[:,1]
    zd = Rd[:,2]
    x =  R[:,0]
    y =  R[:,1]
    z =  R[:,2]
    er = (1/2) * (np.cross(x, xd) + np.cross(y, yd) + np.cross(z, zd))
    assert len(ep.shape) == 1
    assert len(er.shape) == 1
    res = np.hstack((ep, er))
    return res

def ikin(kin, pd, Rd, J, q, iterations=25):
    q_init = q.copy()
    for _ in range(iterations):
        e = totalError(kin, pd, Rd, q)
        kin.Jac(q, J)
        pinv = J.T.dot(np.linalg.inv(J.dot(J.T)))
        delta = pinv.dot(e)
        q += delta
    for i in range(q.shape[0]):
        q[i] = closest_rdr_mod_2pi(q_init[i], q[i])

#
#  Basic Marker
#
def BasicMarker():
    # Configure the basics/shared parameters of the markers.
    marker = Marker()
    marker.header.frame_id    = "world"
    marker.header.stamp       = rospy.Time.now()
    marker.ns                 = "environment"
    marker.action             = Marker.ADD
    marker.pose.orientation.x = 0.0
    marker.pose.orientation.y = 0.0
    marker.pose.orientation.z = 0.0
    marker.pose.orientation.w = 1.0
    marker.color.a            = 1.0     # Make non-transparent!
    return marker

def spline(t0, t1, t):
    if t <= t0: return 0
    if t >= t1: return 1
    x = (t - t0) / (t1 - t0)
    return x * x * (3 - 2 * x)


#
#  Add a Cube
#
def AddCube(markers, position, scale, color):
    # Add a single cube
    markers.append(BasicMarker())
    markers[-1].id   = len(markers)
    markers[-1].type = Marker.CUBE
    markers[-1].pose.position.x = position[0]
    markers[-1].pose.position.y = position[1]
    markers[-1].pose.position.z = position[2]
    markers[-1].scale.x         = scale[0]
    markers[-1].scale.y         = scale[1]
    markers[-1].scale.z         = scale[2]
    markers[-1].color.r         = color[0]
    markers[-1].color.g         = color[1]
    markers[-1].color.b         = color[2]


#
#  Add a Cradle
#
def AddCradle(markers, p, colors):
    # Set the dimensions.
    d = 0.1     # Diameter of cube
    l = 0.06    # Length of rail
    w = 0.02    # Width of rail
    s = (d+w)/2 # Shift from center of cube to center of rail

    # Pull out the colors.
    (cbase, cposx, cnegx, cposy, cnegy) = colors

    # Add the elements.
    AddCube(markers, (p[0],   p[1],   d/2),   (d, d, d), cbase)
    # AddCube(markers, (p[0]+s, p[1],   d/2+s), (w, l, w), cposx)
    # AddCube(markers, (p[0]-s, p[1],   d/2+s), (w, l, w), cnegx)
    # AddCube(markers, (p[0],   p[1]+s, d/2+s), (l, w, w), cposy)
    # AddCube(markers, (p[0],   p[1]-s, d/2+s), (l, w, w), cnegy)

fruit_ids = []
GRAV_ACCEL = -0.5
WORKSPACE_CENTER = np.array([0., 0., 0.6])
WORKSPACE_RADIUS_MIN = 0.2
WORKSPACE_RADIUS_MAX = 0.7
INITIAL_MOVE_TIME = 1
SLICE_BACKWARD = -1.5
SLICE_FORWARD = 0.7
SLICE_TIME = 0.2


# stb
# should satisfy
# interp(-stb, SLICE_TIME - stb, 0)

SLICE_TIME_BEFORE = 0
SLICE_TIME_AFTER = 0

for s in np.arange(0, SLICE_TIME, 0.00001):
    SLICE_TIME_BEFORE = s
    SLICE_TIME_AFTER = SLICE_TIME - s
    i = spline(-SLICE_TIME_BEFORE, SLICE_TIME_AFTER, 0)
    if (1 - i) * SLICE_BACKWARD + i * SLICE_FORWARD >= 0:
        break

class Fruit:
    def __init__(self, ipos, ivel, tzero):
        assert isinstance(ipos, np.ndarray)
        assert isinstance(ivel, np.ndarray)

        self.ipos = np.copy(ipos)
        self.ivel = np.copy(ivel)
        self.tzero = t

        # find next unused id
        i = 0
        while i in fruit_ids:
            i += 1
        fruit_ids.append(i)

        scale = (0.1, 0.1, 0.1)
        color = (1, 0, 1)

        self.marker = BasicMarker()
        self.marker.id = i
        self.marker.type = Marker.CUBE
        self.marker.scale.x         = scale[0]
        self.marker.scale.y         = scale[1]
        self.marker.scale.z         = scale[2]
        self.marker.color.r         = color[0]
        self.marker.color.g         = color[1]
        self.marker.color.b         = color[2]
        self.moveMarker(0)

    def hit(self):
        self.marker.color.r = 1
        self.marker.color.g = 0
        self.marker.color.b = 0

    def pos(self, t):
        ts = t - self.tzero
        return np.array([
            self.ipos[0]  + ts * self.ivel[0],
            self.ipos[1]  + ts * self.ivel[1],
            self.ipos[2]  + ts * self.ivel[2] + (GRAV_ACCEL / 2) * ts ** 2
        ])
    
    def vel(self, t):
        return np.array([
            self.ivel[0],
            self.ivel[1],
            self.ivel[2] + GRAV_ACCEL * (t - self.tzero)
        ])

    def moveMarker(self, t):
        # markers[-1].id   = len(markers)
        # markers[-1].type = Marker.CUBE
        pos = self.pos(t)
        self.marker.pose.position.x = pos[0]
        self.marker.pose.position.y = pos[1]
        self.marker.pose.position.z = pos[2]
        # markers[-1].scale.x         = scale[0]
        # markers[-1].scale.y         = scale[1]
        # markers[-1].scale.z         = scale[2]
        # markers[-1].color.r         = color[0]
        # markers[-1].color.g         = color[1]
        # markers[-1].color.b         = color[2]

    def cleanUp(self):
        fruit_ids.remove(self.marker.id)

SETTING_UP = 0
SLICING = 1

#
#  Main Code
#
if __name__ == "__main__":
    # Prepare the node.
    rospy.init_node('showtouchandgo')

    # Prepare a publisher (latching so new subscribers also get the msg).
    viz_pub = rospy.Publisher("/visualization_marker_array", MarkerArray,
                          queue_size=1, latch=True)


    joint_pub = rospy.Publisher("/joint_states", JointState, queue_size=100)

    urdf = rospy.get_param('/robot_description')
    kin = Kinematics(urdf, 'world', 'tip')
    N   = kin.dofs()
    rospy.loginfo("Loaded URDF for %d joints" % N)


    t = 0
    rate  = 100;
    servo = rospy.Rate(rate)
    dt    = servo.sleep_dur.to_sec()
    rospy.loginfo("Running the servo loop with dt of %f seconds (%fHz)" %
                  (dt, rate))
    q = np.zeros(kin.dofs())
    J = np.zeros((6,kin.dofs())) # Jacobian


    goal_pos = np.array([0.3, 0.3, 0.6])
    goal_rot = np.identity(3)

    setup_init = np.random.uniform(-1, 1, kin.dofs())
    setup_final = np.random.uniform(-1, 1, kin.dofs())
    setup_t0 = 0
    setup_t1 = 1
    col_time = 0

    done_slicing = True

    slice_final = np.random.uniform(-1, 1, kin.dofs())
    currently_slicing = None

    cubes = [ ]
    
    while not rospy.is_shutdown():
        if done_slicing:
            
            new_fruit = Fruit(
                np.array([-3., 0., 0.]),
                np.array([
                    np.random.uniform(0.9, 2),
                    np.random.uniform(-0.4, 0.4),
                    np.random.uniform(0.6, 1.1)
                ]),
                t
            )
            cubes.append(new_fruit)

            # find the time of collision
            tp = t
            collision_times = []
            while True:
                posn = new_fruit.pos(tp)
                if posn[2] < -1:
                    break
                if WORKSPACE_RADIUS_MIN <= np.linalg.norm(WORKSPACE_CENTER - posn) <= WORKSPACE_RADIUS_MAX:
                    collision_times.append(tp)
                tp += dt
            
            if collision_times:
                # choose the first collision that works
                col_time = collision_times[0]
                col_pos  = new_fruit.pos(col_time)
                col_vel  = new_fruit.vel(col_time)
                col_unit_vel = col_vel / np.linalg.norm(col_vel)

                col_rot  = np.identity(3)

                x_rot = -col_vel
                r = col_pos - WORKSPACE_CENTER
                # project direction vector onto x_rot
                z_rot = r -  x_rot * np.dot(r, x_rot) / np.dot(x_rot, x_rot)
                assert np.allclose(0, np.dot(x_rot, z_rot))
                y_rot = np.cross(z_rot, x_rot)

                print(f'Selected a colision at {col_time}')
                print(f'Position: {col_pos}')
                print(f'Velocity: {col_vel}')
                setup_t0 = t
                setup_t1 = min(t + INITIAL_MOVE_TIME, col_time - SLICE_TIME_BEFORE)
                setup_init = slice_final.copy()
                ikin(kin, col_pos, np.vstack((x_rot, y_rot, z_rot)).T, J, setup_final, 500)

                


                slice_final = setup_final.copy()
                slice_final[-1] += SLICE_FORWARD - SLICE_BACKWARD

                # print(f'setup_init = {setup_init}')
                # print(f'setup_final = {setup_final}')

                done_slicing = False
                currently_slicing = new_fruit
            else:
                print(f'Probably cannot slice this fruit')

        t += dt

        new_cubes = []
        for cube in cubes:
            cube.moveMarker(t)
            if cube.pos(t)[2] > -1000:
                new_cubes.append(cube)
            else:
                cube.cleanUp()
        
        cubes.clear()
        cubes.extend(new_cubes)
        
        ar = MarkerArray()
        for cube in cubes:
            ar.markers.append(cube.marker)

        # cube.pose.position.x = np.sin(t)

        viz_pub.publish(ar)

        msg = JointState()
        msg.name.append('theta1')
        msg.name.append('theta2')
        msg.name.append('theta3')
        msg.name.append('theta4')
        msg.name.append('theta5')
        msg.name.append('theta6')

        # if t >= setup_t1:
        #     mode = SETTING_UP
        # elif t <= slice_end_t1

        interp = spline(setup_t0, setup_t1, t)
        joints = interp_joints(setup_init, setup_final, interp)

        swing_interp = spline(col_time - SLICE_TIME_BEFORE, col_time + SLICE_TIME_AFTER, t)
        joints[-1] += (1 - swing_interp) * SLICE_BACKWARD + swing_interp * SLICE_FORWARD
        msg.position = joints
        
        msg.header.stamp = rospy.Time.now()

        if t >= col_time and currently_slicing:
            currently_slicing.hit()
        if t >= col_time + SLICE_TIME_AFTER + 0.5:
            done_slicing = True

        joint_pub.publish(msg)

        servo.sleep()