-
Notifications
You must be signed in to change notification settings - Fork 28
/
test_stacking_script.py
executable file
·175 lines (151 loc) · 8.29 KB
/
test_stacking_script.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
import time
import os
import random
import threading
import argparse
import pdb
from logger import Logger
import matplotlib.pyplot as plt
import numpy as np
import scipy as sc
import cv2
from collections import namedtuple
from robot import Robot
import utils
from utils import StackSequence, annotate_success_manually
from data import GoodRobotDatasetReader
from main import get_and_save_images
from annotate_data import Pair
logger = Logger(True, "scratch", args=None, dir_name="scratch")
############### Testing Block Stacking #######
is_sim = True# Run in simulation?
obj_mesh_dir = os.path.abspath('objects/blocks') if is_sim else None # Directory containing 3D mesh files (.obj) of objects to be added to simulation
log_dir = os.path.abspath('logs/')
logger = Logger(True, log_dir, args=None, dir_name = log_dir)
num_obj = 4 if is_sim else None # Number of objects to add to simulation
tcp_host_ip = args.tcp_host_ip if not is_sim else None # IP and port to robot arm as TCP client (UR5)
tcp_port = args.tcp_port if not is_sim else None
rtc_host_ip = args.rtc_host_ip if not is_sim else None # IP and port to robot arm as real-time client (UR5)
rtc_port = args.rtc_port if not is_sim else None
if is_sim:
workspace_limits = np.asarray([[-0.724, -0.276], [-0.224, 0.224], [-0.0001, 0.4]]) # Cols: min max, Rows: x y z (define workspace limits in robot coordinates)
else:
workspace_limits = np.asarray([[0.3, 0.748], [-0.224, 0.224], [-0.255, -0.1]]) # Cols: min max, Rows: x y z (define workspace limits in robot coordinates)
heightmap_resolution = 0.002 # Meters per pixel of heightmap
random_seed = 1234
force_cpu = False
# -------------- Testing options --------------
is_testing = False
max_test_trials = 1 # Maximum number of test runs per case/scenario
test_preset_cases = False
test_preset_file = os.path.abspath(args.test_preset_file) if test_preset_cases else None
# ------ Pre-loading and logging options ------
#load_snapshot = args.load_snapshot # Load pre-trained snapshot of model?
#snapshot_file = os.path.abspath(args.snapshot_file) if load_snapshot else None
#continue_logging = args.continue_logging # Continue logging from previous session
#logging_directory = os.path.abspath(args.logging_directory) if continue_logging else os.path.abspath('logs')
#save_visualizations = args.save_visualizations # Save visualizations of FCN predictions? Takes 0.6s per training step if set to True
# Set random seed
np.random.seed(random_seed)
# Do we care about color? Switch to
# True to run a color order stacking test,
# False tests stacking order does not matter.
grasp_color_task = False
# are we doing a stack even if we don't care about colors
place_task = True
# place in rows instead of stacks
check_row = False
# if placing in rows, how far apart to set the blocks?
separation = 0.055
distance_threshold = 0.08
robot = Robot(is_sim, obj_mesh_dir, num_obj, workspace_limits,
tcp_host_ip, tcp_port, rtc_host_ip, rtc_port,
is_testing, test_preset_cases, test_preset_file,
place=place_task, grasp_color_task=grasp_color_task)
stacksequence = StackSequence(num_obj, is_goal_conditioned_task=grasp_color_task or place_task)
print('full stack sequence: ' + str(stacksequence.object_color_sequence))
best_rotation_angle = 3.14
blocks_to_move = num_obj - 1
num_stacks = 16
original_position = np.array([-0.6, 0.25, 0])
test_failure = False
valid_depth_heightmap, color_heightmap, depth_heightmap, color_img, depth_img = get_and_save_images(robot, workspace_limits, heightmap_resolution, logger, None, depth_channels_history=False, save_image=False)
prev_heightmap = color_heightmap
next_heightmap = None
dataset_reader_fxn = lambda x: GoodRobotDatasetReader(path_or_obj=x,
split_type="none",
task_type="rows",
augment_by_flipping=False,
augment_language = False,
augment_by_rotating=False,
leave_out_color=None,
batch_size=1,
max_seq_length=40,
resolution = 64,
is_bert = True,
overfit=False)
for stack in range(num_stacks):
print('++++++++++++++++++++++++++++++++++++++++++++++++++')
print('+++++++ Making New Stack ++++++++')
print('++++++++++++++++++++++++++++++++++++++++++++++++++')
theta = 2 * np.pi * stack / num_stacks
# move the first block in order to have it in a standard position.
print('orienting first block')
stack_goal = stacksequence.current_sequence_progress()
block_to_move = stack_goal[0]
block_positions, block_orientations = robot.get_obj_positions_and_orientations()
primitive_position = block_positions[block_to_move]
rotation_angle = block_orientations[block_to_move][2]
robot.grasp(primitive_position, rotation_angle,
object_color=block_to_move)
block_positions = robot.get_obj_positions_and_orientations()[0]
# creates the ideal stack by fixing rotation angle
place = robot.place(original_position.copy(), theta + np.pi / 2)
print('place initial: ' + str(place))
for i in range(blocks_to_move):
print('----------------------------------------------')
stacksequence.next()
stack_goal = stacksequence.current_sequence_progress()
pair = Pair.from_nonsim_main_idxs(color_heightmap,
valid_depth_heightmap,
is_row = check_row)
language_data_instance = dataset_reader_fxn(pair).data['train'][0]
block_to_move = stack_goal[-1]
print('move block: ' + str(i) + ' current stack goal: ' + str(stack_goal))
block_positions, block_orientations = robot.get_obj_positions_and_orientations()
valid_depth_heightmap, color_heightmap, depth_heightmap, color_img, depth_img = get_and_save_images(robot, workspace_limits, heightmap_resolution, logger, None, depth_channels_history=False, save_image=False)
prev_heightmap = color_heightmap
primitive_position = block_positions[block_to_move]
rotation_angle = block_orientations[block_to_move][2]
robot.grasp(primitive_position, rotation_angle, object_color=block_to_move)
base_block_to_place = stack_goal[0]
if check_row:
primitive_position = original_position + (i + 1) * separation * np.array([np.cos(theta), np.sin(theta), 0])
else:
block_positions = robot.get_obj_positions_and_orientations()[0]
primitive_position = block_positions[base_block_to_place]
# place height should be on the top of the stack. otherwise the sim freaks out
stack_z_height = 0
for block_pos in block_positions:
if block_pos[2] > stack_z_height and block_pos[2] < workspace_limits[2][1]:
stack_z_height = block_pos[2]
primitive_position[2] = stack_z_height
place = robot.place(primitive_position, theta + np.pi / 2)
print('place ' + str(i) + ' : ' + str(place))
valid_depth_heightmap, color_heightmap, depth_heightmap, color_img, depth_img = get_and_save_images(robot, workspace_limits, heightmap_resolution, logger, None, depth_channels_history=False, save_image=False)
next_heightmap = color_heightmap
# check if we don't care about color
if not grasp_color_task:
# Deliberately change the goal stack order to test the non-ordered check
stack_goal = np.random.permutation(stack_goal)
print('fake stack goal to test any stack order: ' + str(stack_goal))
if check_row:
stack_success, height_count = robot.check_row(stack_goal, vert_distance_threshold=distance_threshold)
else:
success_code, comment = annotate_success_manually("command", prev_heightmap, next_heightmap)
stack_success, height_count = robot.check_stack(stack_goal, vert_distance_threshold=distance_threshold)
print('stack success part ' + str(i+1) + ' of ' + str(blocks_to_move) + ': ' + str(stack_success))
# reset scene
robot.reposition_objects()
# determine first block to grasp
stacksequence.next()