-
Notifications
You must be signed in to change notification settings - Fork 28
/
generate_sim_stacking_demo.py
executable file
·175 lines (152 loc) · 8.74 KB
/
generate_sim_stacking_demo.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 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
from logger import Logger
def get_and_save_images(stack_num, robot, logger, filename_poststring='0', save_image=True):
# Get latest RGB-D image
valid_depth_heightmap, color_heightmap, depth_heightmap, _, color_img, depth_img = robot.get_camera_data(return_heightmaps=True)
# Save RGB-D images and RGB-D heightmaps
if save_image:
logger.save_images(stack_num, color_img, depth_img, filename_poststring)
logger.save_heightmaps(stack_num, color_heightmap, valid_depth_heightmap, filename_poststring)
return valid_depth_heightmap, color_heightmap, depth_heightmap, color_img, depth_img
if __name__ == '__main__':
####### 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
num_obj = 8 if is_sim else None # Number of objects to add to simulation
final_stack_height = 4 # final desired stack height
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(final_stack_height, 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 = final_stack_height - 1
num_stacks = 3
original_position = np.array([-0.6, 0.25, 0])
# initialize logger
logger = Logger(continue_logging=False, logging_directory='demos')
logger.save_camera_info(robot.cam_intrinsics, robot.cam_pose, robot.cam_depth_scale) # Save camera intrinsics and pose
logger.save_heightmap_info(workspace_limits, heightmap_resolution) # Save heightmap parameters
for stack in range(num_stacks):
print('++++++++++++++++++++++++++++++++++++++++++++++++++')
print('+++++++ Making New Stack ++++++++')
print('++++++++++++++++++++++++++++++++++++++++++++++++++')
theta = 2 * np.pi * stack / num_stacks
original_position = np.append(np.random.uniform(workspace_limits[:2, 0], workspace_limits[:2, 1]), np.zeros(1))
executed_action_log = []
# 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()
#print('\n\n\n')
#print('workspace limits', workspace_limits)
#print('block_positions', block_positions)
#input()
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))
# save initial config
get_and_save_images(stack, robot, workspace_limits, heightmap_resolution, logger, 'orig')
for i in range(blocks_to_move):
print('----------------------------------------------')
stacksequence.next()
stack_goal = stacksequence.current_sequence_progress()
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()
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)
# TODO(adit98) use ACTION_TO_ID from utils.py here
# write action (1 for grasp)
executed_action_log.append(primitive_position + [rotation_angle, 1])
# save post grasp config
get_and_save_images(stack, robot, workspace_limits, heightmap_resolution, logger, str(i) + 'grasp')
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)
# write action (2 for place)
# TODO(adit98) use ACTION_TO_ID from utils.py here
executed_action_log.append(primitive_position + [theta + np.pi / 2, 2])
# save post place config
get_and_save_images(stack, robot, heightmap_resolution, logger, str(i) + 'place')
print('place ' + str(i) + ' : ' + str(place))
# 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, distance_threshold=distance_threshold)
else:
stack_success, height_count = robot.check_stack(stack_goal, distance_threshold=distance_threshold)
print('stack success part ' + str(i+1) + ' of ' + str(blocks_to_move) + ': ' + str(stack_success))
# write actions for finished stack
print('executed-actions-' + str(stack))
logger.write_to_log('executed-actions-' + str(stack), executed_action_log)
# determine first block to grasp
stacksequence.next()