├── .gitignore ├── util.py ├── model.lua ├── README.md ├── LICENSE ├── evaluate.lua ├── visualize.py ├── model-late-concat.lua ├── prepareData.lua ├── evaluateRGB.lua ├── modelRGB.lua ├── dataset.py ├── val.py ├── split.py ├── trainRGB.lua ├── train.lua ├── preprocess.py ├── model.py ├── baxter ├── mirror.py ├── random_sample.py ├── ik_execute.py ├── motion_routine.py ├── execute.py └── servo.py └── train.py /.gitignore: -------------------------------------------------------------------------------- 1 | .DS_Store 2 | *.pyc 3 | processed/ 4 | data/ 5 | splits/ 6 | *.t7 7 | -------------------------------------------------------------------------------- /util.py: -------------------------------------------------------------------------------- 1 | import os 2 | import tensorflow as tf 3 | 4 | def convert_joint_to_dict(joint): 5 | return dict(zip(joint.name, joint.position)) 6 | 7 | def restore_vars(saver, sess, checkpoint_dir): 8 | """ Restore saved net, global score and step, and epsilons OR 9 | create checkpoint directory for later storage. """ 10 | sess.run(tf.initialize_all_variables()) 11 | 12 | if not os.path.exists(checkpoint_dir): 13 | try: 14 | os.makedirs(checkpoint_dir) 15 | except OSError: 16 | pass 17 | 18 | path = tf.train.latest_checkpoint(checkpoint_dir) 19 | if path is None: 20 | print 'no existing checkpoint found' 21 | return False 22 | else: 23 | print 'restoring from %s' % path 24 | saver.restore(sess, path) 25 | return True 26 | -------------------------------------------------------------------------------- /model.lua: -------------------------------------------------------------------------------- 1 | require 'torch' 2 | require 'nn' 3 | 4 | local CROP_SIZE = 128 5 | 6 | local model = nn.Sequential() 7 | 8 | model:add(nn.JoinTable(1, 3)) 9 | model:add(nn.SpatialConvolution(4, 32, 3, 3, 1, 1, 1, 1)) 10 | model:add(nn.ReLU(true)) 11 | 12 | model:add(nn.SpatialConvolution(32, 64, 3, 3, 2, 2, 1, 1)) 13 | model:add(nn.ReLU(true)) 14 | 15 | model:add(nn.SpatialConvolution(64, 128, 3, 3, 2, 2, 1, 1)) 16 | model:add(nn.ReLU(true)) 17 | 18 | model:add(nn.SpatialConvolution(128, 64, 3, 3, 2, 2, 1, 1)) 19 | model:add(nn.ReLU(true)) 20 | 21 | model:add(nn.SpatialConvolution(64, 64, 3, 3, 1, 1, 1, 1)) 22 | model:add(nn.ReLU(true)) 23 | 24 | model:add(nn.SpatialConvolution(64, 64, 3, 3, 1, 1, 1, 1)) 25 | model:add(nn.ReLU(true)) 26 | 27 | model:add(nn.SpatialConvolution(64, 32, 3, 3, 1, 1, 1, 1)) 28 | model:add(nn.ReLU(true)) 29 | 30 | model:add(nn.Reshape(16*16*32)) 31 | model:add(nn.Linear(16*16*32, 2)) 32 | 33 | return model 34 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Robotic grasp 2 | 3 | Use learning to solve the robotic grasp problem. 4 | 5 | ## approach 6 | 7 | We factored the system into 3 components: 8 | 9 | - proposer: take in an image (RGB+D) then propose a grasp to be evaluated 10 | - evaluator: score the grasps 11 | - executor: execute the recommended grasps 12 | 13 | and we will learn/optimize each component using data collected on and off a robot (Baxter). 14 | 15 | ## datasets 16 | 17 | ### training the grasp evaluator 18 | 19 | To train a reasonable grasp evaluator, we treated the evaluation problem as a binary classification task and used the publicly available [Cornell grasp dataset](http://pr.cs.cornell.edu/grasping/rect_data/data.php) to train off the robot. Here are some of statistics of this dataset: 20 | 21 | - image count: 885 22 | - labeled grasps count: 8019 23 | - positive: 5110 (64%) 24 | - negative: 2909 (36%) 25 | - object count: 244 26 | - object category count: 93 27 | 28 | ## usage 29 | 30 | ### use pre-trained models 31 | 32 | ### training 33 | 34 | 35 | ## authors 36 | Falcon Dai (dai@ttic.edu) 37 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2018 Falcon Dai 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /evaluate.lua: -------------------------------------------------------------------------------- 1 | require 'torch' 2 | require 'nn' 3 | require 'cunn' 4 | 5 | local CROP_SIZE = 128 6 | local BATCH_SIZE = arg[2] or 800 7 | 8 | function accuracy (prediction, target) 9 | local _, yHat = torch.max(prediction, 2) 10 | return yHat:eq(target):mean() 11 | end 12 | 13 | local model = torch.load(arg[1]) 14 | local loss = nn.CrossEntropyCriterion() 15 | loss:cuda() 16 | 17 | local val = torch.load('val.t7') 18 | 19 | -- evaluate on val 20 | model:evaluate() 21 | 22 | local nVal = val['y']:size(1) 23 | local totalEnt = 0 24 | local totalCorrect = 0 25 | print('# of val samples:', nVal) 26 | for i = 1, math.ceil(nVal/BATCH_SIZE) do 27 | local j = (i-1) * BATCH_SIZE + 1 28 | local k = math.min(j + BATCH_SIZE - 1, nVal) 29 | local rgb = val['x'][1][{{j, k}}]:cuda() 30 | local d = val['x'][2][{{j, k}}]:cuda() 31 | local y = val['y'][{{j, k}}]:cuda() 32 | local yHat = model:forward({rgb, d}) 33 | 34 | local valCost = loss:forward(yHat, y) 35 | local acc = accuracy(yHat, y) 36 | totalEnt = totalEnt + valCost * (k - j + 1) 37 | totalCorrect = totalCorrect + acc * (k - j + 1) 38 | end 39 | print('val entropy:', totalEnt / nVal) 40 | print('val acc:', totalCorrect / nVal) 41 | -------------------------------------------------------------------------------- /visualize.py: -------------------------------------------------------------------------------- 1 | from PIL import Image, ImageDraw 2 | import argparse 3 | 4 | from dataset import read_label_file, convert_pcd 5 | 6 | parser = argparse.ArgumentParser() 7 | parser.add_argument('partial_path') 8 | 9 | args = parser.parse_args() 10 | 11 | green = (0, 255, 0) 12 | red = (255, 0, 0) 13 | 14 | def draw_box(p1, p2, p3, p4, draw, color): 15 | ps = [p1, p2, p3, p4] 16 | for i in xrange(4): 17 | x1, y1 = ps[i] 18 | x2, y2 = ps[(i+1) % 4] 19 | if i % 2 == 1: 20 | draw.line((x1, y1, x2, y2), fill=color) 21 | else: 22 | draw.line((x1, y1, x2, y2)) 23 | 24 | depth_data = convert_pcd('%s.txt' % args.partial_path) 25 | depth_img = Image.new('F', (640, 480)) 26 | # strange scale in the data 27 | depth_img.putdata(depth_data.flatten() * 1e40) 28 | depth_img.show() 29 | 30 | with Image.open('%sr.png' % args.partial_path) as img: 31 | draw = ImageDraw.Draw(img) 32 | for i, box in enumerate(read_label_file('%scpos.txt' % args.partial_path)): 33 | draw_box(box[0], box[1], box[2], box[3], draw, green) 34 | draw.text(box[0], str(i), fill=green) 35 | 36 | for i, box in enumerate(read_label_file('%scneg.txt' % args.partial_path)): 37 | draw_box(box[0], box[1], box[2], box[3], draw, red) 38 | draw.text(box[0], str(i), fill=red) 39 | 40 | img.show() 41 | -------------------------------------------------------------------------------- /model-late-concat.lua: -------------------------------------------------------------------------------- 1 | require 'torch' 2 | require 'nn' 3 | 4 | local CROP_SIZE = 128 5 | 6 | local model = nn.Sequential() 7 | 8 | local concat = nn.ParallelTable() 9 | local rgb = nn.Sequential() 10 | local d = nn.Sequential() 11 | 12 | rgb:add(nn.SpatialConvolution(3, 32, 3, 3, 1, 1, 1, 1)) 13 | rgb:add(nn.ReLU(true)) 14 | 15 | rgb:add(nn.SpatialConvolution(32, 64, 3, 3, 2, 2, 1, 1)) 16 | rgb:add(nn.ReLU(true)) 17 | 18 | rgb:add(nn.SpatialConvolution(64, 128, 3, 3, 2, 2, 1, 1)) 19 | rgb:add(nn.ReLU(true)) 20 | 21 | rgb:add(nn.SpatialConvolution(128, 64, 3, 3, 2, 2, 1, 1)) 22 | rgb:add(nn.ReLU(true)) 23 | 24 | concat:add(rgb) 25 | 26 | d:add(nn.SpatialConvolution(1, 32, 3, 3, 1, 1, 1, 1)) 27 | d:add(nn.ReLU(true)) 28 | 29 | d:add(nn.SpatialConvolution(32, 64, 3, 3, 2, 2, 1, 1)) 30 | d:add(nn.ReLU(true)) 31 | 32 | d:add(nn.SpatialConvolution(64, 128, 3, 3, 2, 2, 1, 1)) 33 | d:add(nn.ReLU(true)) 34 | 35 | d:add(nn.SpatialConvolution(128, 64, 3, 3, 2, 2, 1, 1)) 36 | d:add(nn.ReLU(true)) 37 | 38 | concat:add(d) 39 | model:add(concat) 40 | model:add(nn.JoinTable(1, 3)) 41 | 42 | model:add(nn.Dropout()) 43 | model:add(nn.SpatialConvolution(128, 64, 3, 3, 1, 1, 1, 1)) 44 | model:add(nn.ReLU(true)) 45 | 46 | model:add(nn.Dropout()) 47 | model:add(nn.SpatialConvolution(64, 64, 3, 3, 1, 1, 1, 1)) 48 | model:add(nn.ReLU(true)) 49 | 50 | model:add(nn.Dropout()) 51 | model:add(nn.Reshape(16*16*64)) 52 | model:add(nn.Linear(16*16*64, 2)) 53 | 54 | return model 55 | -------------------------------------------------------------------------------- /prepareData.lua: -------------------------------------------------------------------------------- 1 | require 'torch' 2 | require 'image' 3 | require 'paths' 4 | npy = require 'npy4th' 5 | 6 | local CROP_SIZE = 128 7 | local DEPTH_RESCALE_FACTOR = 1 / 255 8 | 9 | local trainY = npy.loadnpy('splits/train_y.npy') 10 | local valY = npy.loadnpy('splits/val_y.npy') 11 | 12 | -- lua is 1-index 13 | local trainY = trainY + 1 14 | local valY = valY + 1 15 | 16 | -- load the train set 17 | local trainFn = {} 18 | for line in io.open('splits/train_fn.txt'):lines() do 19 | table.insert(trainFn, line:split(' ')) 20 | end 21 | local nTrain = #trainFn 22 | local trainRgb = torch.Tensor(nTrain, 3, CROP_SIZE, CROP_SIZE) 23 | local trainD = torch.Tensor(nTrain, 1, CROP_SIZE, CROP_SIZE) 24 | for i, row in ipairs(trainFn) do 25 | trainRgb[i] = image.load(row[1]) 26 | trainD[i][1] = npy.loadnpy(row[2]) * DEPTH_RESCALE_FACTOR 27 | end 28 | local trainX = {trainRgb, trainD} 29 | print('train set') 30 | print(trainX) 31 | -- save train set 32 | torch.save('train.t7', { x = trainX, y = trainY }) 33 | 34 | -- load the validation set 35 | local valFn = {} 36 | for line in io.open('splits/val_fn.txt'):lines() do 37 | table.insert(valFn, line:split(' ')) 38 | end 39 | local nVal = #valFn 40 | local valRgb = torch.Tensor(nVal, 3, CROP_SIZE, CROP_SIZE) 41 | local valD = torch.Tensor(nVal, 1, CROP_SIZE, CROP_SIZE) 42 | for i, row in ipairs(valFn) do 43 | valRgb[i] = image.load(row[1]) 44 | valD[i][1] = npy.loadnpy(row[2]) * DEPTH_RESCALE_FACTOR 45 | end 46 | local valX = {valRgb, valD} 47 | print('validation set') 48 | print(valX) 49 | -- save validation set 50 | torch.save('val.t7', { x = valX, y = valY }) 51 | -------------------------------------------------------------------------------- /evaluateRGB.lua: -------------------------------------------------------------------------------- 1 | require 'torch' 2 | require 'nn' 3 | require 'cunn' 4 | 5 | local CROP_SIZE = 128 6 | local BATCH_SIZE = arg[2] or 800 7 | 8 | function accuracy (prediction, target) 9 | local _, yHat = torch.max(prediction, 2) 10 | return yHat:eq(target):mean() 11 | end 12 | 13 | local model = torch.load(arg[1]) 14 | local classLoss = nn.CrossEntropyCriterion():cuda() 15 | local depthLoss = nn.MSECriterion():cuda() 16 | 17 | local val = torch.load('val.t7') 18 | 19 | -- evaluate on val 20 | model:evaluate() 21 | 22 | local nVal = val['y']:size(1) 23 | local totalEnt = 0 24 | local totalCorrect = 0 25 | local totalMSE = 0 26 | print('# of val samples:', nVal) 27 | val['yHat'] = torch.zeros(val['x'][2]:size()) 28 | for i = 1, math.ceil(nVal/BATCH_SIZE) do 29 | local j = (i-1) * BATCH_SIZE + 1 30 | local k = math.min(j + BATCH_SIZE - 1, nVal) 31 | local rgb = val['x'][1][{{j, k}}]:cuda() 32 | local d = val['x'][2][{{j, k}}]:reshape(k-j+1, CROP_SIZE * CROP_SIZE):cuda() 33 | local y = val['y'][{{j, k}}]:cuda() 34 | local yHat = model:forward(rgb) 35 | 36 | val['yHat'][{{j, k}}] = yHat[2]:reshape(k-j+1, CROP_SIZE, CROP_SIZE):float() 37 | 38 | local classCost = classLoss:forward(yHat[1], y) 39 | local depthCost = depthLoss:forward(yHat[2], d) 40 | local acc = accuracy(yHat[1], y) 41 | totalEnt = totalEnt + classCost * (k - j + 1) 42 | totalCorrect = totalCorrect + acc * (k - j + 1) 43 | totalMSE = totalMSE + depthCost * (k - j + 1) 44 | end 45 | torch.save('val-out.t7', val) 46 | print('val entropy:', totalEnt / nVal) 47 | print('val acc:', totalCorrect / nVal) 48 | print('val MSE:', totalMSE / nVal) 49 | -------------------------------------------------------------------------------- /modelRGB.lua: -------------------------------------------------------------------------------- 1 | require 'torch' 2 | require 'nn' 3 | 4 | local CROP_SIZE = 128 5 | 6 | local model = nn.Sequential() 7 | 8 | model:add(nn.SpatialConvolution(3, 32, 3, 3, 1, 1, 1, 1)) 9 | model:add(nn.ReLU(true)) 10 | 11 | model:add(nn.SpatialConvolution(32, 64, 3, 3, 2, 2, 1, 1)) 12 | model:add(nn.ReLU(true)) 13 | 14 | model:add(nn.SpatialConvolution(64, 128, 3, 3, 2, 2, 1, 1)) 15 | model:add(nn.ReLU(true)) 16 | 17 | model:add(nn.SpatialConvolution(128, 64, 3, 3, 2, 2, 1, 1)) 18 | model:add(nn.ReLU(true)) 19 | 20 | model:add(nn.SpatialConvolution(64, 64, 3, 3, 1, 1, 1, 1)) 21 | model:add(nn.ReLU(true)) 22 | 23 | model:add(nn.SpatialConvolution(64, 64, 3, 3, 1, 1, 1, 1)) 24 | model:add(nn.ReLU(true)) 25 | 26 | model:add(nn.SpatialConvolution(64, 32, 3, 3, 1, 1, 1, 1)) 27 | model:add(nn.ReLU(true)) 28 | 29 | local concat = nn.ConcatTable() 30 | local depth = nn.Sequential() 31 | local classify = nn.Sequential() 32 | 33 | classify:add(nn.SpatialConvolution(32, 64, 3, 3, 1, 1, 1, 1)) 34 | classify:add(nn.ReLU(true)) 35 | 36 | classify:add(nn.SpatialConvolution(64, 64, 3, 3, 1, 1, 1, 1)) 37 | classify:add(nn.ReLU(true)) 38 | 39 | classify:add(nn.SpatialConvolution(64, 32, 3, 3, 1, 1, 1, 1)) 40 | classify:add(nn.ReLU(true)) 41 | 42 | classify:add(nn.Reshape(16*16*32)) 43 | classify:add(nn.Linear(16*16*32, 2)) 44 | 45 | concat:add(classify) 46 | 47 | depth:add(nn.SpatialFullConvolution(32, 32, 3, 3, 2, 2, 1, 1, 1, 1)) 48 | depth:add(nn.ReLU(true)) 49 | 50 | depth:add(nn.SpatialFullConvolution(32, 4, 3, 3, 4, 4, 1, 1, 3, 3)) 51 | depth:add(nn.ReLU(true)) 52 | 53 | depth:add(nn.SpatialConvolution(4, 1, 7, 7, 1, 1, 3, 3)) 54 | depth:add(nn.ReLU(true)) 55 | 56 | depth:add(nn.Reshape(CROP_SIZE * CROP_SIZE)) 57 | 58 | concat:add(depth) 59 | 60 | model:add(concat) 61 | 62 | return model 63 | -------------------------------------------------------------------------------- /dataset.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from pylab import * 3 | from PIL import Image, ImageChops 4 | 5 | RAW_WIDTH = 640 6 | RAW_HEIGHT = 480 7 | 8 | def read_label_file(path): 9 | with open(path) as f: 10 | xys = [] 11 | has_nan = False 12 | for l in f: 13 | x, y = map(float, l.split()) 14 | # XXX some bounding boxes has nan coordinates 15 | if np.isnan(x) or np.isnan(y): 16 | has_nan = True 17 | xys.append((x, y)) 18 | if len(xys) % 4 == 0 and len(xys) / 4 >= 1: 19 | if not has_nan: 20 | yield xys[-4], xys[-3], xys[-2], xys[-1] 21 | has_nan = False 22 | 23 | def convert_pcd(path): 24 | with open(path) as f: 25 | # move pass the header 26 | # http://pointclouds.org/documentation/tutorials/pcd_file_format.php 27 | for _ in xrange(11): 28 | f.readline() 29 | pass 30 | im = np.zeros((RAW_HEIGHT, RAW_WIDTH), dtype='f4') 31 | for l in f: 32 | d, i = l.split()[-2:] 33 | d = float(d) 34 | i = int(i) 35 | x = i % RAW_WIDTH 36 | y = i / RAW_WIDTH 37 | im[y, x] = max(0., d) 38 | return im 39 | 40 | def crop_image(img, box, crop_size): 41 | cx, cy = np.mean(box, axis=0) 42 | (x1, y1), (x2, y2) = box[:2] 43 | # center the image to the bounding box 44 | o = ImageChops.offset(img, int(RAW_WIDTH/2-cx), int(RAW_HEIGHT/2-cy)) 45 | # rotate the gripper axis to the x-axis 46 | r = o.rotate(np.rad2deg(np.arctan2(y2 - y1, x2 - x1))) 47 | # crop the image to a fixed size around the bounding box 48 | return r.crop((RAW_WIDTH/2-crop_size/2, RAW_HEIGHT/2-crop_size/2, 49 | RAW_WIDTH/2+crop_size/2, RAW_HEIGHT/2+crop_size/2)) 50 | -------------------------------------------------------------------------------- /val.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | # TODO load from files using queue 3 | import torchfile 4 | 5 | import numpy as np 6 | import time, sys 7 | 8 | from model import build_model 9 | from util import * 10 | 11 | # constants 12 | width = 128 13 | loss_lambda = 0.1 14 | checkpoint_dir = sys.argv[1] 15 | 16 | # model 17 | # grasp_class_prediction, depth_prediction, logit, grasp_image_ph, keep_prob_ph = build_model(width) 18 | grasp_class_prediction, logit, grasp_image_ph, keep_prob_ph = build_model(width) 19 | depth_image_ph = tf.placeholder('float', [None, width, width, 1]) 20 | grasp_class_ph = tf.placeholder('int64', [None]) 21 | 22 | # loss 23 | grasp_class_loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(logit, grasp_class_ph), name='grasp_class_loss') 24 | # depth_loss = tf.reduce_mean(tf.square(depth_image_ph - depth_prediction), name='depth_loss') 25 | # combined_loss = (1. - loss_lambda) * grasp_class_loss + loss_lambda * depth_loss 26 | combined_loss = grasp_class_loss 27 | 28 | # evaluation 29 | batch = int(sys.argv[2]) 30 | correct_prediction = tf.equal(tf.argmax(grasp_class_prediction, 1), grasp_class_ph) 31 | accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32)) 32 | 33 | saver = tf.train.Saver(max_to_keep=5, keep_checkpoint_every_n_hours=1) 34 | with tf.Session() as sess: 35 | restore_vars(saver, sess, checkpoint_dir) 36 | val_data = torchfile.load('val.t7') 37 | n = len(val_data['x'][0]) 38 | print '%d samples' % n 39 | acc = 0. 40 | loss = 0. 41 | for i in xrange(n / batch + 1): 42 | start = batch * i 43 | if n == start: 44 | break 45 | end = min(start + batch, n) 46 | rgb_image = val_data['x'][0][start:end].transpose(0, 2, 3, 1) 47 | grasp_class = val_data['y'][start:end] - 1 48 | eval_feed = { 49 | grasp_image_ph: rgb_image, 50 | grasp_class_ph: grasp_class, 51 | # depth_image_ph: d_image, 52 | keep_prob_ph: 1., 53 | } 54 | 55 | loss += combined_loss.eval(feed_dict=eval_feed) * (end - start) 56 | acc += accuracy.eval(feed_dict=eval_feed) * (end - start) 57 | print acc / n, loss / n 58 | -------------------------------------------------------------------------------- /split.py: -------------------------------------------------------------------------------- 1 | import glob 2 | from pylab import * 3 | import os 4 | from os import path 5 | 6 | def ln(fns, dst_dir): 7 | for p in fns: 8 | fn = path.basename(p) 9 | d = path.dirname(p) 10 | name = fn[:-len('.png')] 11 | os.symlink(p, path.join(dst_dir, fn)) 12 | os.symlink(path.join(d, name + '.npy'), path.join(dst_dir, name + '.npy')) 13 | 14 | np.random.seed(0) 15 | 16 | # create train/val splits 17 | pos_fns = glob.glob('processed/pos/*.png') 18 | neg_fns = glob.glob('processed/neg/*.png') 19 | pos_fns.sort() 20 | neg_fns.sort() 21 | shuffle(pos_fns) 22 | shuffle(neg_fns) 23 | 24 | # using 4:1 train:val ratio as the paper 25 | pos_split = int(floor(len(pos_fns) * 4. / 5.)) 26 | neg_split = int(floor(len(neg_fns) * 4. / 5.)) 27 | 28 | train_pos = pos_fns[:pos_split] 29 | train_neg = neg_fns[:neg_split] 30 | val_pos = pos_fns[pos_split:] 31 | val_neg = neg_fns[neg_split:] 32 | 33 | # set up train/val split folders 34 | os.mkdir('splits') 35 | 36 | os.makedirs('splits/train/pos') 37 | os.makedirs('splits/train/neg') 38 | os.makedirs('splits/val/pos') 39 | os.makedirs('splits/val/neg') 40 | 41 | # make symbolic links 42 | ln(train_pos, 'splits/train/pos') 43 | ln(train_neg, 'splits/train/neg') 44 | ln(val_pos, 'splits/val/pos') 45 | ln(val_neg, 'splits/val/neg') 46 | 47 | # create a randomly permutated list 48 | train_fn = train_pos + train_neg 49 | train_y = [1] * len(train_pos) + [0] * len(train_neg) 50 | ii = np.random.permutation(len(train_fn)) 51 | with open('splits/train_fn.txt', 'wb') as f: 52 | for i in ii: 53 | p = train_fn[i] 54 | name = p[:-len('.png')] 55 | f.write('%s %s\n' % (p, name+'.npy')) 56 | np.save('splits/train_y.npy', asarray(train_y)[ii]) 57 | 58 | val_fn = val_pos + val_neg 59 | val_y = [1] * len(val_pos) + [0] * len(val_neg) 60 | ii = np.random.permutation(len(val_fn)) 61 | with open('splits/val_fn.txt', 'wb') as f: 62 | for i in ii: 63 | p = val_fn[i] 64 | name = p[:-len('.png')] 65 | f.write('%s %s\n' % (p, name+'.npy')) 66 | np.save('splits/val_y.npy', asarray(val_y)[ii]) 67 | 68 | print 'train/val split' 69 | n_train = pos_split + neg_split 70 | print 'train: %i val: %i' % (n_train, len(pos_fns)+len(neg_fns)-n_train) 71 | -------------------------------------------------------------------------------- /trainRGB.lua: -------------------------------------------------------------------------------- 1 | require 'torch' 2 | require 'cutorch' 3 | require 'nn' 4 | require 'cunn' 5 | require 'optim' 6 | require 'paths' 7 | 8 | local BATCH_SIZE = tonumber(arg[2]) or 200 9 | local CROP_SIZE = 128 10 | local MAX_STEP = tonumber(arg[3]) or 800 11 | 12 | function accuracy (prediction, target) 13 | local _, yHat = torch.max(prediction, 2) 14 | return yHat:eq(target):mean() 15 | end 16 | 17 | print('# of CUDA devices:', cutorch.getDeviceCount()) 18 | print('using device:', cutorch.getDevice()) 19 | print('saving checkpoint models to:', arg[1]) 20 | paths.mkdir(arg[1]) 21 | 22 | torch.manualSeed(3) 23 | 24 | local model = require './modelRGB' 25 | print(model) 26 | model:cuda() 27 | model:training() 28 | 29 | local loss = nn.ParallelCriterion() 30 | local classLoss = nn.CrossEntropyCriterion() 31 | local depthLoss = nn.MSECriterion() 32 | local lambda = 0.1 33 | loss:add(classLoss) 34 | loss:add(depthLoss, lambda) 35 | loss:cuda() 36 | 37 | local train = torch.load('train.t7') 38 | 39 | local n = train['y']:size(1) 40 | print('# of samples', n) 41 | 42 | local mParams, mGrad = model:getParameters() 43 | local classCost, depthCost, cost 44 | function _fgrad (rgb, d, y) 45 | function fgrad (params) 46 | mParams:copy(params) 47 | model:zeroGradParameters() 48 | local yHat = model:forward(rgb) 49 | classCost = classLoss:forward(yHat[1], y) 50 | depthCost = depthLoss:forward(yHat[2], d) 51 | cost = loss:forward(yHat, {y, d}) 52 | local dl = loss:backward(yHat, {y, d}) 53 | model:backward(rgb, dl) 54 | return cost, mGrad 55 | end 56 | return fgrad 57 | end 58 | 59 | local rgb, d, y 60 | local state = {} 61 | for step = 1, MAX_STEP do 62 | -- construct mini-batch 63 | local i = step * BATCH_SIZE % n 64 | if i < BATCH_SIZE then 65 | i = 1 66 | end 67 | local j = math.min(i + BATCH_SIZE - 1, n) 68 | local size = j - i + 1 69 | rgb = train['x'][1][{{i, j}}]:cuda() 70 | d = train['x'][2][{{i, j}}]:reshape(size, CROP_SIZE * CROP_SIZE):cuda() 71 | y = train['y'][{{i, j}}]:cuda() 72 | 73 | optim.adam(_fgrad(rgb, d, y), mParams, state) 74 | print(step, cost, classCost, depthCost, mGrad:norm()) 75 | 76 | -- checkpoint the model 77 | if step % 200 == 0 then 78 | model:clearState() 79 | torch.save(arg[1]..'/model.'..step..'.t7', model) 80 | end 81 | end 82 | 83 | -- save the final model 84 | model:clearState() 85 | torch.save(arg[1]..'/model.'..MAX_STEP..'.t7', model) 86 | -------------------------------------------------------------------------------- /train.lua: -------------------------------------------------------------------------------- 1 | require 'torch' 2 | require 'cutorch' 3 | require 'nn' 4 | require 'cunn' 5 | require 'optim' 6 | require 'paths' 7 | 8 | local BATCH_SIZE = tonumber(arg[2]) or 200 9 | local CROP_SIZE = 128 10 | local MAX_STEP = tonumber(arg[3]) or 800 11 | 12 | function accuracy (prediction, target) 13 | local _, yHat = torch.max(prediction, 2) 14 | return yHat:eq(target):mean() 15 | end 16 | 17 | print('# of CUDA devices:', cutorch.getDeviceCount()) 18 | print('using device:', cutorch.getDevice()) 19 | print('saving checkpoint models to:', arg[1]) 20 | paths.mkdir(arg[1]) 21 | 22 | torch.manualSeed(3) 23 | 24 | local model = require './model' 25 | print(model) 26 | model:cuda() 27 | model:training() 28 | local loss = nn.CrossEntropyCriterion() 29 | loss:cuda() 30 | 31 | local train = torch.load('train.t7') 32 | local val = torch.load('val.t7') 33 | 34 | local n = train['y']:size(1) 35 | print('# of samples', n) 36 | 37 | local mParams, mGrad = model:getParameters() 38 | local cost 39 | function _fgrad (rgb, d, y) 40 | function fgrad (params) 41 | mParams:copy(params) 42 | model:zeroGradParameters() 43 | local yHat = model:forward({rgb, d}) 44 | cost = loss:forward(yHat, y) 45 | local dl = loss:backward(yHat, y) 46 | model:backward({rgb, d}, dl) 47 | return cost, mGrad 48 | end 49 | return fgrad 50 | end 51 | 52 | local rgb, d, y 53 | local state = {} 54 | for step = 1, MAX_STEP do 55 | -- construct mini-batch 56 | local i = step * BATCH_SIZE % n 57 | if i < BATCH_SIZE then 58 | i = 1 59 | end 60 | local j = math.min(i + BATCH_SIZE - 1, n) 61 | rgb = train['x'][1][{{i, j}}]:cuda() 62 | d = train['x'][2][{{i, j}}]:cuda() 63 | -- rgb only 64 | -- d:zero() 65 | -- depth only 66 | rgb:zero() 67 | y = train['y'][{{i, j}}]:cuda() 68 | 69 | optim.adam(_fgrad(rgb, d, y), mParams, state) 70 | print(step, cost, mGrad:norm()) 71 | 72 | -- checkpoint the model 73 | if step % 200 == 0 then 74 | model:clearState() 75 | torch.save(arg[1]..'/model.'..step..'.t7', model) 76 | end 77 | end 78 | 79 | -- evaluate on val 80 | model:evaluate() 81 | 82 | rgb = val['x'][1][{{1,1000}}]:cuda() 83 | d = val['x'][2][{{1,1000}}]:cuda() 84 | y = val['y'][{{1,1000}}]:cuda() 85 | local yHat = model:forward({rgb, d}) 86 | 87 | local valCost = loss:forward(yHat, y) 88 | print('val entropy:', valCost) 89 | print('val acc:', accuracy(yHat, y)) 90 | 91 | model:clearState() 92 | torch.save(arg[1]..'/model.'..MAX_STEP..'.t7', model) 93 | -------------------------------------------------------------------------------- /preprocess.py: -------------------------------------------------------------------------------- 1 | from PIL import Image 2 | import argparse 3 | import glob 4 | import os 5 | 6 | from dataset import * 7 | 8 | CROP_SIZE = 128 9 | # strange scale in the depth data 10 | DEPTH_SCALE_FACTOR = 1e40 11 | 12 | parser = argparse.ArgumentParser() 13 | parser.add_argument('dataset_path') 14 | parser.add_argument('processed_dataset_path') 15 | 16 | args = parser.parse_args() 17 | 18 | # set up folders 19 | try: 20 | os.mkdir(args.processed_dataset_path) 21 | except: 22 | pass 23 | try: 24 | os.mkdir('%s/pos' % args.processed_dataset_path) 25 | os.mkdir('%s/neg' % args.processed_dataset_path) 26 | except: 27 | pass 28 | 29 | obj_cat = {} 30 | with open('%s/processedData/z.txt' % args.dataset_path) as description_f: 31 | for line in description_f: 32 | sid, obj_id, category = line.split()[:3] 33 | obj_cat[sid] = (obj_id, category) 34 | 35 | # file format string 36 | # /-. 37 | filename_format = '%s-%03i' 38 | 39 | dimg = Image.new('F', (RAW_WIDTH, RAW_HEIGHT)) 40 | n_img = 0 41 | n_pos = 0 42 | n_neg = 0 43 | objs = set() 44 | cats = set() 45 | for path in glob.glob('%s/*/pcd*[0-9].txt' % args.dataset_path): 46 | print path 47 | n_img += 1 48 | sample_id = path[-len('1234.txt'):-len('.txt')] 49 | objs.add(obj_cat[sample_id][0]) 50 | cats.add(obj_cat[sample_id][1]) 51 | dim = convert_pcd(path) 52 | # replace NaN with 0 53 | dimg.putdata(np.nan_to_num(dim.flatten() * DEPTH_SCALE_FACTOR)) 54 | with Image.open(path[:-len('.txt')]+'r.png') as cimg: 55 | # positive grasps 56 | for i, box in enumerate(read_label_file(path[:-len('.txt')]+'cpos.txt')): 57 | n_pos += 1 58 | filename = filename_format % (sample_id, i) 59 | crop_image(cimg, box, CROP_SIZE).save('%s/pos/%s.png' % (args.processed_dataset_path, filename)) 60 | np.save('%s/pos/%s.npy' % (args.processed_dataset_path, filename), np.reshape(crop_image(dimg, box, CROP_SIZE).getdata(), (CROP_SIZE, CROP_SIZE))) 61 | 62 | # negative grasps 63 | for i, box in enumerate(read_label_file(path[:-len('.txt')]+'cneg.txt')): 64 | n_neg += 1 65 | filename = filename_format % (sample_id, i) 66 | crop_image(cimg, box, CROP_SIZE).save('%s/neg/%s.png' % (args.processed_dataset_path, filename)) 67 | np.save('%s/neg/%s.npy' % (args.processed_dataset_path, filename), np.reshape(crop_image(dimg, box, CROP_SIZE).getdata(), (CROP_SIZE, CROP_SIZE))) 68 | 69 | n_grasp = n_pos + n_neg 70 | print 71 | print 'dataset statistics:' 72 | print '# of objects:', len(objs) 73 | print '# of object categories:', len(cats) 74 | print '# of images:', n_img 75 | print '# of labeled grasps: %i positive: %i (%.2f) negative: %i (%.2f)' % (n_grasp, n_pos, n_pos * 1./n_grasp, n_neg, n_neg * 1./n_grasp) 76 | -------------------------------------------------------------------------------- /model.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | 4 | def build_model(width): 5 | with tf.name_scope('model'): 6 | grasp_image_ph = tf.placeholder('float', [None, width, width, 3]) 7 | keep_prob_ph = tf.placeholder('float', name='dropout') 8 | 9 | # rgb conv 10 | a1 = tf.contrib.layers.convolution2d(tf.nn.dropout(grasp_image_ph, keep_prob_ph), 32, (5, 5), tf.nn.relu, weight_init=tf.contrib.layers.xavier_initializer_conv2d(), name='conv1') 11 | # a2 = tf.contrib.layers.convolution2d(tf.nn.dropout(a1, keep_prob_ph), 64, (3, 3), tf.nn.relu, stride=(2, 2), weight_init=tf.contrib.layers.xavier_initializer_conv2d(), name='conv2') 12 | # a3 = tf.contrib.layers.convolution2d(tf.nn.dropout(a2, keep_prob_ph), 128, (3, 3), tf.nn.relu, stride=(2, 2), weight_init=tf.contrib.layers.xavier_initializer_conv2d(), name='conv3') 13 | # a4 = tf.contrib.layers.convolution2d(tf.nn.dropout(a3, keep_prob_ph), 64, (3, 3), tf.nn.relu, stride=(2, 2), weight_init=tf.contrib.layers.xavier_initializer_conv2d(), name='conv4') 14 | 15 | a2 = tf.contrib.layers.convolution2d(tf.nn.dropout(a1, keep_prob_ph), 64, (3, 3), tf.nn.relu, weight_init=tf.contrib.layers.xavier_initializer_conv2d(), name='conv2') 16 | a2_max = tf.nn.max_pool(a2, (1, 3, 3, 1), (1, 2, 2, 1), 'SAME') 17 | a3 = tf.contrib.layers.convolution2d(tf.nn.dropout(a2_max, keep_prob_ph), 64, (3, 3), tf.nn.relu, weight_init=tf.contrib.layers.xavier_initializer_conv2d(), name='conv3') 18 | a3_max = tf.nn.max_pool(a3, (1, 3, 3, 1), (1, 2, 2, 1), 'SAME') 19 | a4 = tf.contrib.layers.convolution2d(tf.nn.dropout(a3_max, keep_prob_ph), 32, (3, 3), tf.nn.relu, weight_init=tf.contrib.layers.xavier_initializer_conv2d(), name='conv4') 20 | a4_max = tf.nn.max_pool(a4, (1, 3, 3, 1), (1, 2, 2, 1), 'SAME') 21 | 22 | conv = a4_max 23 | 24 | # flatten 25 | conv_shape = conv.get_shape().as_list() 26 | flat_dim = np.product(conv_shape[1:]) 27 | print 'final shape', conv_shape, 'flat_dim', flat_dim 28 | conv_flat = tf.reshape(conv, [-1, flat_dim]) 29 | 30 | # fc 31 | fc1 = tf.contrib.layers.fully_connected(conv_flat, 2, weight_init=tf.contrib.layers.xavier_initializer(), name='fc1') 32 | 33 | # depth 34 | # with tf.variable_scope('depth'): 35 | # batch_size = tf.shape(grasp_image_ph)[0] 36 | # deconv_shape1 = tf.pack([batch_size, 64, 64, 16]) 37 | # f1 = tf.Variable(tf.contrib.layers.xavier_initializer_conv2d()([5, 5, 16, 64]), name='features1') 38 | # d1 = tf.nn.conv2d_transpose(a3, f1, deconv_shape1, (1, 1, 1, 1), name='deconv1') 39 | # deconv_shape2 = tf.pack([batch_size, width, width, 1]) 40 | # f2 = tf.Variable(tf.contrib.layers.xavier_initializer_conv2d()([5, 5, 1, 16]), name='features2') 41 | # d2 = tf.nn.conv2d_transpose(tf.nn.relu(d1), f2, deconv_shape2, (1, 2, 2, 1), name='deconv2') 42 | 43 | # prediction 44 | logit = fc1 45 | grasp_class_prediction = tf.nn.softmax(fc1) 46 | # depth_prediction = d2 47 | 48 | # return grasp_class_prediction, depth_prediction, logit, grasp_image_ph, keep_prob_ph 49 | return grasp_class_prediction, logit, grasp_image_ph, keep_prob_ph 50 | -------------------------------------------------------------------------------- /baxter/mirror.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | import sys 4 | 5 | import rospy 6 | 7 | import baxter_interface 8 | 9 | from baxter_interface import CHECK_VERSION 10 | from baxter_core_msgs.msg import DigitalIOState, EndEffectorState 11 | from sensor_msgs.msg import JointState, Range, Image 12 | 13 | from PIL import Image as PIL_Image 14 | from PIL import ImageDraw as PIL_ImageDraw 15 | import numpy as np 16 | 17 | def main(): 18 | # parser = argparse.ArgumentParser() 19 | # parser.add_argument('controlling_limb', default='right', choices=['left', 'right']) 20 | # args = parser.parse_args() 21 | 22 | rospy.init_node('robot_mirror_control') 23 | 24 | left = baxter_interface.Limb('left') 25 | right = baxter_interface.Limb('right') 26 | grip_left = baxter_interface.Gripper('left', CHECK_VERSION) 27 | grip_right = baxter_interface.Gripper('right', CHECK_VERSION) 28 | head = baxter_interface.Head() 29 | pub = rospy.Publisher('/robot/xdisplay', Image, queue_size=1) 30 | measurements = { 31 | 'ir_range': 0., 32 | 'grip_force': 0., 33 | 'grip_gap': 0., 34 | } 35 | 36 | def close_grippers(data): 37 | if data.state == DigitalIOState.PRESSED: 38 | grip_right.close() 39 | grip_left.close() 40 | 41 | def open_grippers(data): 42 | if data.state == DigitalIOState.PRESSED: 43 | grip_right.open() 44 | grip_left.open() 45 | 46 | def change_left_limb(data): 47 | x = {} 48 | right_arm_joint_names = ['right_e0', 'right_e1', 'right_s0', 'right_s1', 'right_w0', 'right_w1', 'right_w2'] 49 | for name, position in zip(data.name, data.position): 50 | if name in right_arm_joint_names: 51 | x[name.replace('right', 'left')] = position 52 | if 'left_s0' in x: 53 | HEAD_LIMB_OFFSET = -0.76 54 | head.set_pan(x['left_s0'] + HEAD_LIMB_OFFSET) 55 | x['left_s0'] *= -1 56 | if len(x.keys()) > 0: 57 | left.set_joint_positions(x) 58 | 59 | def display_gripper_state(data): 60 | # print 'position:', data.position, 'force:', data.force 61 | measurements['grip_force'] = data.force 62 | measurements['grip_gap'] = data.position 63 | 64 | def display_ir(data): 65 | # print 'left range:', data.range 66 | measurements['ir_range'] = data.range 67 | 68 | def display_cam(data): 69 | img = PIL_Image.frombytes('RGBA', (data.width, data.height), data.data) 70 | draw = PIL_ImageDraw.Draw(img) 71 | draw.multiline_text((100, 100), 'distance: %.4f\ngrip gap: %.2f\ngrip force: %.4f\n' % (measurements['ir_range'], measurements['grip_gap'], measurements['grip_force'])) 72 | data.data = img.tobytes() 73 | pub.publish(data) 74 | 75 | rs = baxter_interface.RobotEnable(CHECK_VERSION) 76 | try: 77 | rs.enable() 78 | grip_left.calibrate() 79 | grip_right.calibrate() 80 | rospy.Subscriber('/robot/digital_io/right_lower_button/state', DigitalIOState, close_grippers) 81 | rospy.Subscriber('/robot/digital_io/right_upper_button/state', DigitalIOState, open_grippers) 82 | rospy.Subscriber('/robot/joint_states', JointState, change_left_limb) 83 | rospy.Subscriber('/robot/end_effector/left_gripper/state', EndEffectorState, display_gripper_state) 84 | rospy.Subscriber('/robot/range/left_hand_range/state', Range, display_ir) 85 | rospy.Subscriber('/cameras/left_hand_camera/image', Image, display_cam) 86 | 87 | rospy.spin() 88 | except Exception, e: 89 | rospy.logerr(e.strerror) 90 | 91 | return 0 92 | 93 | if __name__ == '__main__': 94 | sys.exit(main()) 95 | -------------------------------------------------------------------------------- /baxter/random_sample.py: -------------------------------------------------------------------------------- 1 | import rospy 2 | import baxter_interface 3 | import numpy as np 4 | 5 | from std_msgs.msg import Header 6 | from baxter_core_msgs.msg import DigitalIOState, EndEffectorState 7 | from sensor_msgs.msg import Image 8 | 9 | from PIL import Image as PIL_Image 10 | from PIL import ImageDraw as PIL_ImageDraw 11 | 12 | from motion_routine import * 13 | from model import build_model 14 | from util import * 15 | 16 | def main(): 17 | rospy.init_node('execute') 18 | rs = baxter_interface.RobotEnable(baxter_interface.CHECK_VERSION) 19 | rs.enable() 20 | 21 | # retrieve images 22 | global current_image 23 | def update_image(msg): 24 | global current_image 25 | current_image = PIL_Image.frombytes('RGBA', (msg.width, msg.height), msg.data) 26 | # print msg.width, msg.height, msg.is_bigendian, msg.step, msg.encoding 27 | rospy.Subscriber('/cameras/left_hand_camera/image', Image, update_image) 28 | 29 | # model 30 | width = 128 31 | checkpoint_dir = 'checkpoints-dev-rgb-4-max' 32 | grasp_class_prediction, logit, grasp_image_ph, keep_prob_ph = build_model(width) 33 | saver = tf.train.Saver(max_to_keep=5, keep_checkpoint_every_n_hours=1) 34 | 35 | arm = baxter_interface.Limb('left') 36 | arm.move_to_neutral() 37 | gripper = baxter_interface.Gripper('left') 38 | gripper.calibrate() 39 | 40 | # grasp crop 41 | crop_center_x = 330 42 | crop_center_y = 160 43 | grasp_class_threashold = 0.5 44 | scale = 1.0 45 | crop_width = width * scale 46 | crop_box = (crop_center_x - crop_width/2, crop_center_y - crop_width/2, crop_center_x + crop_width/2, crop_center_y + crop_width/2) 47 | 48 | # grasp workspace 49 | x0 = 0.81 50 | y0 = 0.25 51 | delta = 0.04 52 | initial_z = 0.1 53 | bound_z = -0.165 54 | 55 | grasp_class_threashold = 0.5 56 | 57 | pub = rospy.Publisher('/robot/xdisplay', Image, queue_size=1) 58 | global force 59 | def display_gripper_state(msg): 60 | global force 61 | force = msg.force 62 | 63 | rospy.Subscriber('/robot/end_effector/left_gripper/state', EndEffectorState, display_gripper_state) 64 | 65 | with tf.Session() as sess: 66 | restore_vars(saver, sess, checkpoint_dir) 67 | 68 | attemp = 0 69 | while True: 70 | # sample a grasp 71 | dx = np.random.rand() * (2. * delta) - delta 72 | dy = np.random.rand() * (2. * delta) - delta 73 | target_theta = (np.random.rand() * 2. - 1.) * 3.059 74 | target_x = x0 + dx 75 | target_y = y0 + dy 76 | 77 | # move to the grasp location 78 | execute_linear(arm, target_x, target_y, initial_z, target_theta) 79 | 80 | # predict grasp 81 | crop = np.array(current_image.crop(crop_box).resize((width, width)))[:,:,:3] 82 | grasp_pred = grasp_class_prediction.eval(session=sess, feed_dict={ 83 | grasp_image_ph: crop.reshape((1, width, width, 3)), 84 | keep_prob_ph: 1., 85 | }) 86 | 87 | # display image 88 | draw = PIL_ImageDraw.Draw(current_image) 89 | draw.text(crop_box[:2], 'prob: %.5f' % grasp_pred[0, 1]) 90 | draw.text((20, 20), 'grasp force: %.5f' % force) 91 | if grasp_pred[0, 1] > grasp_class_threashold: 92 | draw.rectangle(crop_box, outline=(0, 255, 0)) 93 | else: 94 | draw.rectangle(crop_box, outline=(0, 0, 255)) 95 | msg = Image( 96 | header=Header( 97 | stamp=rospy.Time.now(), 98 | frame_id='base', 99 | ), 100 | width=640, 101 | height=400, 102 | step=640 * 4, 103 | encoding='bgra8', 104 | is_bigendian=0, 105 | data=current_image.tobytes(), 106 | ) 107 | pub.publish(msg) 108 | if grasp_pred[0, 1] > grasp_class_threashold: 109 | execute_planar_grasp(arm, gripper, initial_z, bound_z, target_theta, lower_to_drop=0.05) 110 | 111 | attemp += 1 112 | 113 | if __name__ == '__main__': 114 | main() 115 | -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | # TODO load from files using queue 3 | import torchfile 4 | 5 | import numpy as np 6 | import time, sys 7 | 8 | from model import build_model 9 | from util import * 10 | 11 | # constants 12 | width = 128 13 | loss_lambda = 0.1 14 | checkpoint_dir = sys.argv[1] 15 | 16 | # model 17 | # grasp_class_prediction, depth_prediction, logit, grasp_image_ph, keep_prob_ph = build_model(width) 18 | grasp_class_prediction, logit, grasp_image_ph, keep_prob_ph = build_model(width) 19 | depth_image_ph = tf.placeholder('float', [None, width, width, 1]) 20 | grasp_class_ph = tf.placeholder('int64', [None]) 21 | 22 | # loss 23 | grasp_class_loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(logit, grasp_class_ph), name='grasp_class_loss') 24 | # depth_loss = tf.reduce_mean(tf.square(depth_image_ph - depth_prediction), name='depth_loss') 25 | # combined_loss = (1. - loss_lambda) * grasp_class_loss + loss_lambda * depth_loss 26 | combined_loss = grasp_class_loss 27 | 28 | # optimization 29 | batch = int(sys.argv[2]) 30 | n_eval_interval = int(sys.argv[3]) 31 | n_train_step = 10**3 32 | global_step = tf.Variable(0, trainable=False, name='global_step') 33 | initial_learning_rate = 0.001 34 | decay_steps = 64 35 | decay_rate = 0.9 36 | momentum = 0.5 37 | learning_rate = tf.train.exponential_decay(initial_learning_rate, global_step, decay_steps, decay_rate, staircase=True) 38 | train_op = tf.train.AdamOptimizer(learning_rate).minimize(combined_loss, global_step=global_step) 39 | # train_op = tf.train.GradientDescentOptimizer(learning_rate).minimize(grasp_class_loss, global_step=global_step) 40 | # train_op = tf.train.MomentumOptimizer(learning_rate, momentum).minimize(grasp_class_loss, global_step=global_step) 41 | 42 | # evaluation 43 | correct_prediction = tf.equal(tf.argmax(grasp_class_prediction, 1), grasp_class_ph) 44 | accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32)) 45 | 46 | # summary 47 | tf.scalar_summary('learning_rate', learning_rate) 48 | tf.scalar_summary('grasp_loss', grasp_class_loss) 49 | # tf.scalar_summary('depth_loss', depth_loss) 50 | tf.scalar_summary('loss', combined_loss) 51 | tf.scalar_summary('accuracy', accuracy) 52 | summary_op = tf.merge_all_summaries() 53 | 54 | def main(): 55 | saver = tf.train.Saver(max_to_keep=5, keep_checkpoint_every_n_hours=1) 56 | with tf.Session() as sess: 57 | tf.set_random_seed(1234) 58 | np.random.seed(123) 59 | 60 | writer = tf.train.SummaryWriter('tf-log/%d' % time.time(), sess.graph_def) 61 | 62 | restore_vars(saver, sess, checkpoint_dir) 63 | 64 | # load train data 65 | train_data = torchfile.load('train.t7') 66 | n = len(train_data['x'][0]) 67 | for i in xrange(n_train_step): 68 | # start = i * batch % n 69 | # end = min(start + batch, n) 70 | # rgb_image = train_data['x'][0][start:end].transpose(0, 2, 3, 1) 71 | # grasp_class = train_data['y'][start:end] - 1 72 | 73 | ind = np.random.choice(n, batch, replace=False) 74 | rgb_image = train_data['x'][0][ind].transpose(0, 2, 3, 1) 75 | # d_image = train_data['x'][1][ind].transpose(0, 2, 3, 1) 76 | grasp_class = train_data['y'][ind] - 1 77 | 78 | if i % n_eval_interval == 0: 79 | val_feed = { 80 | grasp_image_ph: rgb_image, 81 | grasp_class_ph: grasp_class, 82 | # depth_image_ph: d_image, 83 | keep_prob_ph: 1., 84 | } 85 | 86 | # print logit.eval(feed_dict=val_feed) 87 | print 'grasp loss', grasp_class_loss.eval(feed_dict=val_feed) 88 | # print 'depth loss', depth_loss.eval(feed_dict=val_feed) 89 | print 'accuracy', accuracy.eval(feed_dict=val_feed) 90 | writer.add_summary(sess.run(summary_op, feed_dict=val_feed), i) 91 | saver.save(sess, checkpoint_dir + '/model', global_step=i) 92 | 93 | # train 94 | train_feed = { 95 | grasp_image_ph: rgb_image, 96 | grasp_class_ph: grasp_class, 97 | # depth_image_ph: d_image, 98 | keep_prob_ph: 0.8, 99 | } 100 | train_op.run(feed_dict=train_feed) 101 | 102 | # save the model 103 | saver.save(sess, checkpoint_dir + '/model', global_step=i) 104 | 105 | if __name__ == '__main__': 106 | main() 107 | -------------------------------------------------------------------------------- /baxter/ik_execute.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | import sys 4 | 5 | import rospy 6 | 7 | import baxter_interface 8 | import actionlib 9 | 10 | from baxter_interface import CHECK_VERSION 11 | from geometry_msgs.msg import ( 12 | PoseStamped, 13 | Pose, 14 | Point, 15 | Quaternion, 16 | ) 17 | from std_msgs.msg import Header 18 | 19 | from baxter_core_msgs.srv import ( 20 | SolvePositionIK, 21 | SolvePositionIKRequest, 22 | ) 23 | from baxter_core_msgs.msg import DigitalIOState 24 | from sensor_msgs.msg import JointState, Range, Image 25 | from control_msgs.msg import FollowJointTrajectoryAction, FollowJointTrajectoryGoal 26 | from trajectory_msgs.msg import JointTrajectoryPoint 27 | 28 | # from PIL import Image as PIL_Image 29 | # from PIL import ImageDraw as PIL_ImageDraw 30 | import numpy as np 31 | import cv2 32 | from cv_bridge import CvBridge 33 | 34 | class Trajectory(object): 35 | def __init__(self, limb): 36 | ns = 'robot/limb/' + limb + '/' 37 | self._client = actionlib.SimpleActionClient( 38 | ns + "follow_joint_trajectory", 39 | FollowJointTrajectoryAction, 40 | ) 41 | self._goal = FollowJointTrajectoryGoal() 42 | self._goal_time_tolerance = rospy.Time(0.1) 43 | self._goal.goal_time_tolerance = self._goal_time_tolerance 44 | server_up = self._client.wait_for_server(timeout=rospy.Duration(10.0)) 45 | if not server_up: 46 | rospy.logerr("Timed out waiting for Joint Trajectory" 47 | " Action Server to connect. Start the action server" 48 | " before running example.") 49 | rospy.signal_shutdown("Timed out waiting for Action Server") 50 | sys.exit(1) 51 | self.clear(limb) 52 | 53 | def add_point(self, positions, time): 54 | point = JointTrajectoryPoint() 55 | point.positions = positions 56 | point.time_from_start = rospy.Duration(time) 57 | self._goal.trajectory.points.append(point) 58 | 59 | def start(self): 60 | self._goal.trajectory.header.stamp = rospy.Time.now() 61 | self._client.send_goal(self._goal) 62 | 63 | def stop(self): 64 | self._client.cancel_goal() 65 | rospy.sleep(0.1) 66 | 67 | def wait(self, timeout=15.0): 68 | self._client.wait_for_result(timeout=rospy.Duration(timeout)) 69 | 70 | def result(self): 71 | return self._client.get_result() 72 | 73 | def clear(self, limb): 74 | self._goal = FollowJointTrajectoryGoal() 75 | self._goal.goal_time_tolerance = self._goal_time_tolerance 76 | self._goal.trajectory.joint_names = [limb + '_' + joint for joint in \ 77 | ['s0', 's1', 'e0', 'e1', 'w0', 'w1', 'w2']] 78 | 79 | def get_ik_joints(target_x, target_y, initial_z, target_z, target_theta, n_steps): 80 | ns = "ExternalTools/left/PositionKinematicsNode/IKService" 81 | iksvc = rospy.ServiceProxy(ns, SolvePositionIK) 82 | ikreq = SolvePositionIKRequest(seed_mode=SolvePositionIKRequest.SEED_CURRENT) 83 | hdr = Header(stamp=rospy.Time.now(), frame_id='base') 84 | 85 | qx = np.sin(target_theta * 0.5) 86 | qy = np.cos(target_theta * 0.5) 87 | 88 | for z in np.arange(initial_z, target_z, (target_z-initial_z)/n_steps): 89 | pose = PoseStamped( 90 | header=hdr, 91 | pose=Pose( 92 | position=Point( x=target_x, y=target_y, z=z, ), 93 | orientation=Quaternion( x=qx, y=qy, z=0., w=0., ), 94 | ), 95 | ) 96 | ikreq.pose_stamp.append(pose) 97 | try: 98 | rospy.wait_for_service(ns, 5.0) 99 | resp = iksvc(ikreq) 100 | except (rospy.ServiceException, rospy.ROSException), e: 101 | rospy.logerr("Service call failed: %s" % (e,)) 102 | return 1 103 | 104 | return [j for v, j in zip(resp.isValid, resp.joints) if v] 105 | 106 | def planar_grasp(arm, gripper, target_x, target_y, initial_z, target_z, target_theta, grasp_range_threshold=0.13, n_steps=20): 107 | # build trajectory 108 | traj = Trajectory('left') 109 | rospy.on_shutdown(traj.stop) 110 | current_angles = [arm.joint_angle(joint) for joint in arm.joint_names()] 111 | traj.add_point(current_angles, 0.0) 112 | 113 | for i, joint in enumerate(get_ik_joints(target_x, target_y, initial_z, target_z, target_theta, n_steps)): 114 | traj.add_point(joint.position, 1. + 0.1 * i) 115 | 116 | global sub 117 | def near_object(msg): 118 | # print 'range', msg.range 119 | if msg.range < grasp_range_threshold: 120 | global sub 121 | sub.unregister() 122 | print 'near object' 123 | traj.stop() 124 | gripper.close() 125 | # wait for the gripper to close 126 | rospy.sleep(1.) 127 | 128 | # lift arm 129 | traj2 = Trajectory('left') 130 | s = arm.joint_angles() 131 | current_z = arm.endpoint_pose()['position'].z 132 | ss = [s[x] for x in traj2._goal.trajectory.joint_names] 133 | traj2.add_point(ss, 0.) 134 | for j, joint in enumerate(get_ik_joints(target_x, target_y, current_z, initial_z, target_theta, 10)): 135 | traj2.add_point(joint.position, 0.2 * j + 1.) 136 | traj2.start() 137 | 138 | sub = rospy.Subscriber('/robot/range/left_hand_range/state', Range, near_object) 139 | 140 | # move arm 141 | traj.start() 142 | return traj 143 | 144 | def main(): 145 | rospy.init_node('execute') 146 | rs = baxter_interface.RobotEnable(CHECK_VERSION) 147 | rs.enable() 148 | 149 | arm = baxter_interface.Limb('left') 150 | gripper = baxter_interface.Gripper('left') 151 | 152 | gripper.calibrate() 153 | 154 | def open_grippers(msg): 155 | if msg.state == DigitalIOState.PRESSED: 156 | gripper.open() 157 | 158 | rospy.Subscriber('/robot/digital_io/left_upper_button/state', DigitalIOState, open_grippers) 159 | 160 | # target position 161 | target_x = 0.7 162 | target_y = 0.3 163 | initial_z = 0.1 164 | target_z = -0.17 165 | target_theta = np.pi / 4 166 | 167 | planar_grasp(arm, gripper, target_x, target_y, initial_z, target_z, target_theta) 168 | rospy.spin() 169 | print 'executed trajectory' 170 | 171 | return 0 172 | 173 | 174 | if __name__ == '__main__': 175 | main() 176 | -------------------------------------------------------------------------------- /baxter/motion_routine.py: -------------------------------------------------------------------------------- 1 | import sys 2 | 3 | import rospy 4 | import actionlib 5 | from std_msgs.msg import Header 6 | from geometry_msgs.msg import ( 7 | PoseStamped, 8 | Pose, 9 | Point, 10 | Quaternion, 11 | ) 12 | from baxter_core_msgs.srv import ( 13 | SolvePositionIK, 14 | SolvePositionIKRequest, 15 | ) 16 | from control_msgs.msg import ( 17 | FollowJointTrajectoryAction, 18 | FollowJointTrajectoryGoal, 19 | ) 20 | from trajectory_msgs.msg import JointTrajectoryPoint 21 | from util import convert_joint_to_dict 22 | 23 | class Trajectory(object): 24 | def __init__(self, limb): 25 | ns = 'robot/limb/' + limb + '/' 26 | self._client = actionlib.SimpleActionClient( 27 | ns + "follow_joint_trajectory", 28 | FollowJointTrajectoryAction, 29 | ) 30 | self._goal = FollowJointTrajectoryGoal() 31 | self._goal_time_tolerance = rospy.Time(0.1) 32 | self._goal.goal_time_tolerance = self._goal_time_tolerance 33 | server_up = self._client.wait_for_server(timeout=rospy.Duration(10.0)) 34 | if not server_up: 35 | rospy.logerr("Timed out waiting for Joint Trajectory" 36 | " Action Server to connect. Start the action server" 37 | " before running example.") 38 | rospy.signal_shutdown("Timed out waiting for Action Server") 39 | sys.exit(1) 40 | self.clear(limb) 41 | 42 | def add_point(self, positions, time): 43 | point = JointTrajectoryPoint() 44 | point.positions = positions 45 | point.time_from_start = rospy.Duration(time) 46 | self._goal.trajectory.points.append(point) 47 | 48 | def start(self): 49 | self._goal.trajectory.header.stamp = rospy.Time.now() 50 | self._client.send_goal(self._goal) 51 | 52 | def stop(self): 53 | self._client.cancel_goal() 54 | rospy.sleep(0.1) 55 | 56 | def wait(self, timeout=15.0): 57 | self._client.wait_for_result(timeout=rospy.Duration(timeout)) 58 | 59 | def result(self): 60 | return self._client.get_result() 61 | 62 | def clear(self, limb): 63 | self._goal = FollowJointTrajectoryGoal() 64 | self._goal.goal_time_tolerance = self._goal_time_tolerance 65 | self._goal.trajectory.joint_names = [limb + '_' + joint for joint in \ 66 | ['s0', 's1', 'e0', 'e1', 'w0', 'w1', 'w2']] 67 | 68 | def get_ik_joints_linear(initial_x, initial_y, initial_z, initial_w2, 69 | target_x, target_y, target_z, target_w2, n_steps): 70 | ns = "ExternalTools/left/PositionKinematicsNode/IKService" 71 | iksvc = rospy.ServiceProxy(ns, SolvePositionIK) 72 | ikreq = SolvePositionIKRequest(seed_mode=SolvePositionIKRequest.SEED_CURRENT) 73 | hdr = Header(stamp=rospy.Time.now(), frame_id='base') 74 | 75 | # current_pose = arm.endpoint_pose() 76 | x0 = initial_x 77 | y0 = initial_y 78 | z0 = initial_z 79 | 80 | # linear interpolate between current pose and target pose 81 | for i in xrange(n_steps): 82 | t = (i + 1) * 1. / n_steps 83 | x = (1. - t) * x0 + t * target_x 84 | y = (1. - t) * y0 + t * target_y 85 | z = (1. - t) * z0 + t * target_z 86 | 87 | pose = PoseStamped( 88 | header=hdr, 89 | pose=Pose( 90 | position=Point( x=x, y=y, z=z, ), 91 | # endeffector pointing down 92 | orientation=Quaternion( x=1., y=0., z=0., w=0., ), 93 | ), 94 | ) 95 | ikreq.pose_stamp.append(pose) 96 | try: 97 | rospy.wait_for_service(ns, 5.0) 98 | resp = iksvc(ikreq) 99 | except (rospy.ServiceException, rospy.ROSException), e: 100 | rospy.logerr("Service call failed: %s" % (e,)) 101 | return [] 102 | 103 | js = [] 104 | # control w2 separately from other joints 105 | for i, (v, j) in enumerate(zip(resp.isValid, resp.joints)): 106 | t = (i + 1) * 1. / n_steps 107 | if v: 108 | w2 = (1. - t) * initial_w2 + t * target_w2 109 | j.position = j.position[:-1] + (w2,) 110 | js.append(j) 111 | return js 112 | 113 | def get_linear_trajectory_from_current(arm, target_x, target_y, target_z, target_w2, n_steps, duration): 114 | traj = Trajectory('left') 115 | x = arm.joint_angles() 116 | current_pose = arm.endpoint_pose() 117 | current_position = current_pose['position'] 118 | w2i = x['left_w2'] 119 | dt_step = duration * 1. / n_steps 120 | for j, joint in enumerate(get_ik_joints_linear(current_position.x, current_position.y, current_position.z, w2i, target_x, target_y, target_z, target_w2, n_steps)): 121 | traj.add_point(joint.position, dt_step * (j + 1)) 122 | return traj 123 | 124 | def execute_linear(arm, target_x, target_y, target_z, target_w2, n_steps=10, duration=4., timeout=10.): 125 | initial_w2 = arm.joint_angles()['left_w2'] 126 | pos = arm.endpoint_pose()['position'] 127 | traj = get_linear_trajectory_from_current(arm, target_x, target_y, target_z, target_w2, n_steps, duration) 128 | traj.start() 129 | traj.wait(timeout) 130 | return traj 131 | 132 | def execute_vertical(arm, target_z, n_steps=10, duration=4., timeout=10.): 133 | initial_w2 = arm.joint_angles()['left_w2'] 134 | pos = arm.endpoint_pose()['position'] 135 | return execute_linear(arm, pos.x, pos.y, target_z, initial_w2, n_steps, duration, timeout) 136 | 137 | def execute_horizontal(arm, target_x, target_y, target_w2, n_steps=10, duration=4., timeout=10.): 138 | initial_w2 = arm.joint_angles()['left_w2'] 139 | pos = arm.endpoint_pose()['position'] 140 | return execute_linear(arm, target_x, target_y, pos.z, target_w2, n_steps, duration, timeout) 141 | 142 | def execute_planar_grasp(arm, gripper, initial_z, target_z, target_w2, n_steps=10, duration=4., timeout=10., sleep=1., lower_to_drop=0.): 143 | execute_vertical(arm, target_z, n_steps, duration, timeout) 144 | gripper.close() 145 | rospy.sleep(sleep) 146 | execute_vertical(arm, initial_z, n_steps, duration, timeout) 147 | if lower_to_drop > 0: 148 | execute_vertical(arm, target_z + lower_to_drop, n_steps, duration, timeout) 149 | gripper.open() 150 | rospy.sleep(sleep) 151 | 152 | if __name__ == '__main__': 153 | x0 = 0.81 154 | y0 = 0.25 155 | delta = 0.04 156 | initial_z = 0.1 157 | bound_z = -0.165 158 | 159 | for _ in xrange(10): 160 | execute_linear(arm, x0, y0, initial_z, 0.) 161 | execute_planar_grasp(arm, gripper, initial_z, bound_z, 0., lower_to_drop=0.05) 162 | -------------------------------------------------------------------------------- /baxter/execute.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | import sys 4 | 5 | import rospy 6 | 7 | import baxter_interface 8 | import actionlib 9 | 10 | from baxter_interface import CHECK_VERSION 11 | from baxter_core_msgs.msg import DigitalIOState, EndEffectorState 12 | from sensor_msgs.msg import JointState, Range, Image 13 | from control_msgs.msg import FollowJointTrajectoryAction, FollowJointTrajectoryGoal 14 | from trajectory_msgs.msg import JointTrajectoryPoint 15 | 16 | from PIL import Image as PIL_Image 17 | from PIL import ImageDraw as PIL_ImageDraw 18 | import numpy as np 19 | 20 | # baxter's arm links in mm 21 | SHOULDER_LINK = 370.82 22 | ELBOW_LINK = 374.42 23 | SPEED = 100. 24 | 25 | def d_arccos(x): 26 | return -1. / np.sqrt(1 - x**2) 27 | 28 | def angle_dot(r, z, v=SPEED, a=SHOULDER_LINK, b=ELBOW_LINK): 29 | l = np.sqrt(r**2 + z**2) 30 | q = (a**2 + l**2 - b**2) / (2 * a * l) 31 | print l, q 32 | return d_arccos(q) * (l**2 - a**2 + b**2) / (2 * a * l**2) * z / np.sqrt(r**2 + z**2) * v 33 | 34 | def solve_joint(r, z, a=SHOULDER_LINK, b=ELBOW_LINK): 35 | l = np.sqrt(r**2 + z**2) 36 | theta = -np.arctan2(z, r) 37 | print theta 38 | alpha = np.arccos((a**2 + l**2 - b**2) / (2 * a * l)) 39 | beta = alpha + np.arccos((a**2 + l**2 - b**2) / (2 * a * l)) 40 | return -alpha + theta, beta, alpha - theta - beta + np.pi/2 41 | 42 | def main(): 43 | # parser = argparse.ArgumentParser() 44 | # parser.add_argument('controlling_limb', default='right', choices=['left', 'right']) 45 | # args = parser.parse_args() 46 | 47 | rospy.init_node('robot_mirror_control') 48 | 49 | left = baxter_interface.Limb('left') 50 | right = baxter_interface.Limb('right') 51 | grip_left = baxter_interface.Gripper('left', CHECK_VERSION) 52 | grip_right = baxter_interface.Gripper('right', CHECK_VERSION) 53 | head = baxter_interface.Head() 54 | pub = rospy.Publisher('/robot/xdisplay', Image, queue_size=1) 55 | measurements = { 56 | 'ir_range': 0., 57 | 'grip_force': 0., 58 | 'grip_gap': 0., 59 | } 60 | 61 | def close_grippers(data): 62 | if data.state == DigitalIOState.PRESSED: 63 | grip_right.close() 64 | grip_left.close() 65 | 66 | def open_grippers(data): 67 | if data.state == DigitalIOState.PRESSED: 68 | grip_right.open() 69 | grip_left.open() 70 | 71 | def move_limb(r1, z1): 72 | def velocity_control_callback(data): 73 | joints = {} 74 | for name, position in zip(data.name, data.position): 75 | joints[name] = position 76 | print joints 77 | if 'left_s1' in joints and 'left_e1' in joints: 78 | alpha = joints['left_s1'] 79 | beta = joints['left_e1'] 80 | z = SHOULDER_LINK * np.sin(-alpha) - ELBOW_LINK * np.sin(beta) 81 | l = np.sqrt(z**2 + r1**2) 82 | theta = np.arccos(r1 / l) 83 | print alpha, beta, l, theta, z 84 | theta_dot = -1 / (1. + (r1/l)**2) * z / np.sqrt(r1**2 + z**2) * SPEED 85 | alpha_dot = angle_dot(r1, z) 86 | beta_dot = alpha_dot + angle_dot(r1, z, a=ELBOW_LINK, b=SHOULDER_LINK) 87 | left.set_joint_velocities({ 88 | 'left_s1': -alpha_dot + theta_dot, 89 | 'left_e1': -beta_dot, 90 | 'left_w1': alpha_dot - theta_dot - beta_dot 91 | }) 92 | return velocity_control_callback 93 | 94 | def display_gripper_state(data): 95 | # print 'position:', data.position, 'force:', data.force 96 | measurements['grip_force'] = data.force 97 | measurements['grip_gap'] = data.position 98 | 99 | def display_ir(data): 100 | # print 'left range:', data.range 101 | measurements['ir_range'] = data.range 102 | 103 | def display_cam(data): 104 | img = PIL_Image.frombytes('RGBA', (data.width, data.height), data.data) 105 | draw = PIL_ImageDraw.Draw(img) 106 | draw.multiline_text((100, 100), 'distance: %.4f\ngrip gap: %.2f\ngrip force: %.4f\n' % (measurements['ir_range'], measurements['grip_gap'], measurements['grip_force'])) 107 | data.data = img.tobytes() 108 | pub.publish(data) 109 | 110 | rs = baxter_interface.RobotEnable(CHECK_VERSION) 111 | try: 112 | rs.enable() 113 | # grip_left.calibrate() 114 | # grip_right.calibrate() 115 | # rospy.Subscriber('/robot/digital_io/right_lower_button/state', DigitalIOState, close_grippers) 116 | # rospy.Subscriber('/robot/digital_io/right_upper_button/state', DigitalIOState, open_grippers) 117 | rospy.Subscriber('/robot/end_effector/left_gripper/state', EndEffectorState, display_gripper_state) 118 | rospy.Subscriber('/robot/range/left_hand_range/state', Range, display_ir) 119 | rospy.Subscriber('/cameras/left_hand_camera/image', Image, display_cam) 120 | a0, b0, c0 = solve_joint(400., 0.) 121 | lowest = np.sqrt((SHOULDER_LINK + ELBOW_LINK) ** 2 - 500 ** 2) 122 | a1, b1, c1 = solve_joint(400., -lowest) 123 | print a0, b0, c0 124 | print a1, b1, c1 125 | 126 | action_client = actionlib.SimpleActionClient( 127 | '/robot/limb/left/follow_joint_trajectory', 128 | FollowJointTrajectoryAction 129 | ) 130 | server_up = action_client.wait_for_server(timeout=rospy.Duration(10.0)) 131 | if not server_up: 132 | rospy.logerr("Timed out waiting for Joint Trajectory" 133 | " Action Server to connect. Start the action server" 134 | " before running example.") 135 | rospy.signal_shutdown("Timed out waiting for Action Server") 136 | return 1 137 | action_client.cancel_goal() 138 | # left.move_to_neutral() 139 | 140 | goal = FollowJointTrajectoryGoal() 141 | goal.trajectory.joint_names = ['left_s1', 'left_e1', 'left_w1', 'left_s0', 'left_e0', 'left_w0', 'left_w2'] 142 | goal.goal_time_tolerance = rospy.Time(1.) 143 | 144 | # specify trajectory 145 | current_p = JointTrajectoryPoint() 146 | current_p.positions = [left.joint_angle(joint) for joint in goal.trajectory.joint_names] 147 | current_p.time_from_start = rospy.Duration(1) 148 | p0 = JointTrajectoryPoint() 149 | p0.positions = [a0, b0, c0, 0., 0., 0., 0.] 150 | p0.time_from_start = rospy.Duration(3.) 151 | p1 = JointTrajectoryPoint() 152 | p1.positions = [a1, b1, c1, 0., 0., 0., 0.] 153 | p1.time_from_start = rospy.Duration(10.) 154 | goal.trajectory.points.append(current_p) 155 | goal.trajectory.points.append(p0) 156 | goal.trajectory.points.append(p1) 157 | 158 | # left.move_to_joint_positions({ 159 | # 'left_s1': a0, 160 | # 'left_e1': b0, 161 | # 'left_w1': c0, 162 | # }) 163 | # rospy.sleep(2.) 164 | global stopped 165 | stopped = False 166 | def stop_when_near(data): 167 | global stopped 168 | if data.range < 0.2 and not stopped: 169 | print 'reached threashold. stopping arm' 170 | action_client.cancel_goal() 171 | stopped = True 172 | 173 | rospy.Subscriber('/robot/range/left_hand_range/state', Range, stop_when_near) 174 | goal.trajectory.header.stamp = rospy.Time.now() 175 | action_client.send_goal(goal) 176 | rospy.sleep(2.) 177 | if not action_client.wait_for_result(rospy.Duration(20.)): 178 | rospy.logdebug("Canceling goal") 179 | action_client.cancel_goal() 180 | if action_client.wait_for_result(rospy.Duration(20.)): 181 | rospy.logdebug('Preempted') 182 | else: 183 | rospy.logdebug("Preempt didn't finish") 184 | print action_client.get_result() 185 | 186 | # rospy.Subscriber('/robot/joint_states', JointState, move_limb(500., lowest)) 187 | # left.move_to_joint_positions({ 188 | # 'left_s1': a1, 189 | # 'left_e1': b1, 190 | # 'left_w1': c1, 191 | # }) 192 | # rospy.spin() 193 | except Exception, e: 194 | print e 195 | 196 | return 0 197 | 198 | if __name__ == '__main__': 199 | sys.exit(main()) 200 | -------------------------------------------------------------------------------- /baxter/servo.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | from ik_execute import * 4 | from model import build_model 5 | from util import * 6 | 7 | from sensor_msgs.msg import Image 8 | from baxter_core_msgs.msg import DigitalIOState, EndEffectorState 9 | from PIL import Image as PIL_Image 10 | from PIL import ImageDraw as PIL_ImageDraw 11 | 12 | import cv2 13 | from cv_bridge import CvBridge 14 | 15 | import sys 16 | 17 | global current_traj 18 | 19 | def lift_arm(arm, z0, original_theta, n_steps=10): 20 | global current_traj 21 | # if current_traj: 22 | # current_traj.stop() 23 | 24 | traj2 = Trajectory('left') 25 | s = arm.joint_angles() 26 | current_pose = arm.endpoint_pose() 27 | current_position = current_pose['position'] 28 | # theta0 = np.arcsin(current_pose['orientation'].x) 29 | theta0 = original_theta 30 | ss = [s[x] for x in traj2._goal.trajectory.joint_names] 31 | traj2.add_point(ss, 0.) 32 | for j, joint in enumerate(get_ik_joints_linear(arm, current_position.x, current_position.y, current_position.z, theta0, n_steps)): 33 | traj2.add_point(joint.position, 0.2 * j + 1.) 34 | traj2.start() 35 | return traj2 36 | 37 | def ik_solve(x, y, z, z_theta): 38 | ns = "ExternalTools/left/PositionKinematicsNode/IKService" 39 | iksvc = rospy.ServiceProxy(ns, SolvePositionIK) 40 | ikreq = SolvePositionIKRequest(seed_mode=SolvePositionIKRequest.SEED_CURRENT) 41 | hdr = Header(stamp=rospy.Time.now(), frame_id='base') 42 | 43 | qx = np.sin(z_theta * 0.5) 44 | qy = np.cos(z_theta * 0.5) 45 | 46 | pose = PoseStamped( 47 | header=hdr, 48 | pose=Pose( 49 | position=Point( x=x, y=y, z=z, ), 50 | orientation=Quaternion( x=qx, y=qy, z=0., w=0., ), 51 | ), 52 | ) 53 | ikreq.pose_stamp.append(pose) 54 | try: 55 | rospy.wait_for_service(ns, 5.0) 56 | resp = iksvc(ikreq) 57 | except (rospy.ServiceException, rospy.ROSException), e: 58 | rospy.logerr("Service call failed: %s" % (e,)) 59 | return 1 60 | 61 | return resp.joints[0] 62 | 63 | def get_ik_joints_linear(arm, target_x, target_y, target_z, target_theta, n_steps): 64 | ns = "ExternalTools/left/PositionKinematicsNode/IKService" 65 | iksvc = rospy.ServiceProxy(ns, SolvePositionIK) 66 | ikreq = SolvePositionIKRequest(seed_mode=SolvePositionIKRequest.SEED_CURRENT) 67 | hdr = Header(stamp=rospy.Time.now(), frame_id='base') 68 | 69 | current_pose = arm.endpoint_pose() 70 | x0 = current_pose['position'].x 71 | y0 = current_pose['position'].y 72 | z0 = current_pose['position'].z 73 | 74 | try: 75 | w2f = convert_joint_to_dict(ik_solve(target_x, target_y, target_z, target_theta))['left_w2'] 76 | w2i = arm.joint_angles()['left_w2'] 77 | except: 78 | return [] 79 | 80 | # linear interpolate between current pose and target pose 81 | for i in xrange(n_steps): 82 | t = (i + 1) * 1. / n_steps 83 | x = (1. - t) * x0 + t * target_x 84 | y = (1. - t) * y0 + t * target_y 85 | z = (1. - t) * z0 + t * target_z 86 | 87 | pose = PoseStamped( 88 | header=hdr, 89 | pose=Pose( 90 | position=Point( x=x, y=y, z=z, ), 91 | orientation=Quaternion( x=1., y=0., z=0., w=0., ), 92 | ), 93 | ) 94 | ikreq.pose_stamp.append(pose) 95 | try: 96 | rospy.wait_for_service(ns, 5.0) 97 | resp = iksvc(ikreq) 98 | except (rospy.ServiceException, rospy.ROSException), e: 99 | rospy.logerr("Service call failed: %s" % (e,)) 100 | return 1 101 | 102 | js = [] 103 | for i, (v, j) in enumerate(zip(resp.isValid, resp.joints)): 104 | t = (i + 1) * 1. / n_steps 105 | if v: 106 | w2 = (1. - t) * w2i + t * w2f 107 | j.position = j.position[:-1] + (w2,) 108 | js.append(j) 109 | return js 110 | 111 | def planar_grasp(arm, gripper, target_x, target_y, initial_z, target_z, target_theta, grasp_range_threshold=0.12, n_steps=20): 112 | # build trajectory 113 | traj = Trajectory('left') 114 | # rospy.on_shutdown(traj.stop) 115 | # current_angles = [arm.joint_angle(joint) for joint in arm.joint_names()] 116 | # traj.add_point(current_angles, 0.0) 117 | 118 | for i, joint in enumerate(get_ik_joints_linear(arm, target_x, target_y, target_z, target_theta, n_steps)): 119 | traj.add_point(joint.position, 0.2 + 0.3 * i) 120 | 121 | global sub 122 | global current_traj 123 | def near_object(msg): 124 | # print 'range', msg.range 125 | if msg.range < grasp_range_threshold: 126 | global sub 127 | sub.unregister() 128 | print 'near object' 129 | traj.stop() 130 | gripper.close() 131 | # wait for the gripper to close 132 | rospy.sleep(1.) 133 | 134 | # lift arm 135 | traj2 = Trajectory('left') 136 | s = arm.joint_angles() 137 | current_z = arm.endpoint_pose()['position'].z 138 | ss = [s[x] for x in traj2._goal.trajectory.joint_names] 139 | traj2.add_point(ss, 0.) 140 | for j, joint in enumerate(get_ik_joints_linear(arm, target_x, target_y, initial_z, target_theta, n_steps)): 141 | traj2.add_point(joint.position, 0.2 + 0.4 * j) 142 | traj2.start() 143 | current_traj = traj2 144 | traj2.wait(20.) 145 | rospy.sleep(1.) 146 | print 'lifted gripper' 147 | 148 | # if sub == None: 149 | # sub = rospy.Subscriber('/robot/range/left_hand_range/state', Range, near_object) 150 | 151 | # move arm 152 | traj.start() 153 | current_traj = traj 154 | traj.wait(20.) 155 | 156 | # close gripper 157 | gripper.close() 158 | # wait for the gripper to close 159 | rospy.sleep(1.) 160 | 161 | # lift arm 162 | traj2 = Trajectory('left') 163 | s = arm.joint_angles() 164 | current_z = arm.endpoint_pose()['position'].z 165 | ss = [s[x] for x in traj2._goal.trajectory.joint_names] 166 | traj2.add_point(ss, 0.) 167 | for j, joint in enumerate(get_ik_joints_linear(arm, target_x, target_y, initial_z, target_theta, n_steps)): 168 | traj2.add_point(joint.position, 0.2 + 0.2 * j) 169 | traj2.start() 170 | current_traj = traj2 171 | traj2.wait(20.) 172 | rospy.sleep(2.) 173 | print 'lifted gripper' 174 | 175 | return traj, sub 176 | 177 | def main(): 178 | # model 179 | width = 128 180 | checkpoint_dir = 'checkpoints-dev-rgb-4-max' 181 | grasp_class_prediction, logit, grasp_image_ph, keep_prob_ph = build_model(width) 182 | saver = tf.train.Saver(max_to_keep=5, keep_checkpoint_every_n_hours=1) 183 | 184 | with tf.Session() as sess: 185 | restore_vars(saver, sess, checkpoint_dir) 186 | 187 | rospy.init_node('execute') 188 | rs = baxter_interface.RobotEnable(CHECK_VERSION) 189 | rs.enable() 190 | 191 | arm = baxter_interface.Limb('left') 192 | gripper = baxter_interface.Gripper('left') 193 | 194 | # arm.move_to_neutral() 195 | gripper.calibrate() 196 | 197 | def open_grippers(msg): 198 | if msg.state == DigitalIOState.PRESSED: 199 | gripper.open() 200 | 201 | # camera listener 202 | # bridge = CvBridge() 203 | global check_grasp 204 | global found_grasp 205 | found_grasp = False 206 | pub = rospy.Publisher('/robot/xdisplay', Image, queue_size=1) 207 | crop_center_x = 330 208 | crop_center_y = 160 209 | grasp_class_threashold = 0.5 210 | scale = 1.0 211 | crop_width = width * scale 212 | crop_box = (crop_center_x - crop_width/2, crop_center_y - crop_width/2, crop_center_x + crop_width/2, crop_center_y + crop_width/2) 213 | def classify_image(msg): 214 | img = PIL_Image.frombytes('RGBA', (msg.width, msg.height), msg.data) 215 | # img = bridge.imgmsg_to_cv2(msg, desired_encoding="passthrough") 216 | # cv2.imshow('hello', img) 217 | # msg.data = img.crop((msg.width - width/2, msg.height - width/2, msg.width + width/2, msg.height + width/2)).tobytes() 218 | # print np.array(img.crop(crop_box)).shape 219 | crop = np.array(img.crop(crop_box).resize((width, width)))[:,:,:3] 220 | grasp_pred = grasp_class_prediction.eval(session=sess, feed_dict={ 221 | grasp_image_ph: crop.reshape((1, width, width, 3)), 222 | keep_prob_ph: 1., 223 | }) 224 | draw = PIL_ImageDraw.Draw(img) 225 | # print grasp_pred[0, 1] 226 | global force 227 | draw.text(crop_box[:2], 'prob: %.5f' % grasp_pred[0, 1]) 228 | draw.text((20, 20), 'grasp force: %.5f' % force) 229 | if grasp_pred[0, 1] > grasp_class_threashold: 230 | draw.rectangle(crop_box, outline=(0, 255, 0)) 231 | else: 232 | draw.rectangle(crop_box, outline=(0, 0, 255)) 233 | msg.data = img.tobytes() 234 | pub.publish(msg) 235 | 236 | global check_grasp 237 | global found_grasp 238 | global best_p 239 | global best_x 240 | global best_y 241 | global best_theta 242 | 243 | if check_grasp and grasp_pred[0, 1] > grasp_class_threashold: 244 | found_grasp = True 245 | check_grasp = False 246 | best_p = grasp_pred[0, 1] 247 | best_x = target_x 248 | best_y = target_y 249 | best_theta = target_theta 250 | # global current_traj 251 | # current_traj.stop() 252 | # current_traj.wait(10.) 253 | 254 | # traj, sub = planar_grasp(arm, gripper, target_x, target_y, initial_z, bound_z, target_theta) 255 | # # current_traj = traj 256 | # print 'attempting planar grasp' 257 | # current_traj = traj 258 | # traj.wait(20.) 259 | 260 | # best grasp 261 | global best_p 262 | global best_x 263 | global best_y 264 | global best_theta 265 | best_p = 0. 266 | best_x = None 267 | best_y = None 268 | best_theta = None 269 | 270 | rospy.Subscriber('/robot/digital_io/left_upper_button/state', DigitalIOState, open_grippers) 271 | rospy.Subscriber('/cameras/left_hand_camera/image', Image, classify_image) 272 | 273 | global force 274 | def display_gripper_state(msg): 275 | global force 276 | force = msg.force 277 | 278 | rospy.Subscriber('/robot/end_effector/left_gripper/state', EndEffectorState, display_gripper_state) 279 | 280 | # positions 281 | x0 = 0.81 282 | y0 = 0.25 283 | delta = 0.04 284 | initial_z = 0.1 285 | bound_z = -0.17 286 | # n_attempts = int(sys.argv[1]) 287 | 288 | while True: 289 | gripper.open() 290 | found_grasp = False 291 | global current_traj 292 | current_traj = None 293 | traj = None 294 | global sub 295 | sub = None 296 | # for attempt in xrange(n_attempts): 297 | attempt = 0 298 | while not found_grasp: 299 | attempt += 1 300 | check_grasp = False 301 | # if found_grasp: 302 | # print 'found grasp' 303 | # break 304 | # sample a grasp 305 | dx = np.random.rand() * (2. * delta) - delta 306 | dy = np.random.rand() * (2. * delta) - delta 307 | target_theta = (np.random.rand() * 2. - 1.) * 3.059 308 | target_x = x0 + dx 309 | target_y = y0 + dy 310 | # target_x = x0 311 | # target_y = y0 312 | # target_theta = 0. 313 | joint = ik_solve(target_x, target_y, initial_z, target_theta) 314 | 315 | print attempt 316 | print dx, dy, target_theta 317 | print joint.position 318 | 319 | if sub: 320 | sub.unregister() 321 | 322 | if len(joint.position) > 0: 323 | # arm.move_to_joint_positions(dict(zip(joint.name, joint.position))) 324 | traj2 = Trajectory('left') 325 | s = arm.joint_angles() 326 | current_position = arm.endpoint_pose()['position'] 327 | ss = [s[x] for x in traj2._goal.trajectory.joint_names] 328 | traj2.add_point(ss, 0.) 329 | for j, joint in enumerate(get_ik_joints_linear(arm, target_x, target_y, initial_z, target_theta, 10)): 330 | traj2.add_point(joint.position, 0.2 * j + 1.) 331 | traj2.start() 332 | current_traj = traj2 333 | traj2.wait(20.) 334 | print 'moved to position' 335 | rospy.sleep(1.) 336 | check_grasp = True 337 | rospy.sleep(1.) 338 | 339 | print 'best grasp', best_p, best_x, best_y, best_theta 340 | if best_x != None: 341 | print 'executing best grasp' 342 | traj2 = Trajectory('left') 343 | for j, joint in enumerate(get_ik_joints_linear(arm, best_x, best_y, initial_z, best_theta, 10)): 344 | traj2.add_point(joint.position, 0.2 * j + 1.) 345 | traj2.start() 346 | traj2.wait(30.) 347 | print 'moved to best grasp' 348 | # traj2.stop() 349 | print 'attempting planar grasp' 350 | traj, sub = planar_grasp(arm, gripper, best_x, best_y, initial_z, bound_z, best_theta) 351 | # current_traj = traj 352 | current_traj = traj 353 | traj.wait(20.) 354 | 355 | print 'done' 356 | # current_traj = lift_arm(arm, initial_z, target_theta) 357 | # rospy.spin() 358 | # current_traj.wait(10.) 359 | print 'executed trajectory' 360 | 361 | return 0 362 | 363 | 364 | if __name__ == '__main__': 365 | main() 366 | --------------------------------------------------------------------------------