robot.print_info()
# define useful variables for FK
link_id = robot.get_end_effector_ids(end_effector=0)
joint_ids = robot.joints # actuated joint
# load data
with open(os.path.dirname(os.path.abspath(__file__)) + '/data.txt', 'rb') as f:
positions = pickle.load(f)
# set initial joint position
robot.reset_joint_states(q=positions[0])
# draw a sphere at the position of the end-effector
sphere = world.load_visual_sphere(
position=robot.get_link_world_positions(link_id),
radius=0.05,
color=(1, 0, 0, 0.5))
sphere = Body(sim, body_id=sphere)
# perform simulation
for t in count():
# if no more joint positions, get out of the loop
if t >= len(positions):
break
# set joint positions
robot.set_joint_positions(positions[t], joint_ids=joint_ids)
# make the sphere follow the end effector
sphere.position = robot.get_link_world_positions(link_id)
# define useful variables for IK
dt = 1. / 240
link_id = robot.get_end_effector_ids(end_effector=0)
joint_ids = robot.joints # actuated joint
damping = 0.01 # for damped-least-squares IK
wrt_link_id = -1 # robot.get_link_ids('iiwa_link_1')
qIdx = robot.get_q_indices(joint_ids)
# define gains
kp = 50 # 5 if velocity control, 50 if position control
kd = 0 # 2*np.sqrt(kp)
# create sphere to follow
sphere = world.load_visual_sphere(position=np.array([0.5, 0., 1.]),
radius=0.05,
color=(1, 0, 0, 0.5))
sphere = Body(sim, body_id=sphere)
# set initial joint positions (based on the position of the sphere at [0.5, 0, 1])
robot.reset_joint_states(q=[
8.84305270e-05, 7.11378917e-02, -1.68059886e-04, -9.71690439e-01,
1.68308810e-05, 3.71467111e-01, 5.62890805e-05
])
# define amplitude and angular velocity when moving the sphere
w = 0.01
r = 0.2
for t in count():
# move sphere
# create robot
robot = KukaIIWA(sim)
robot.print_info()
# define useful variables for IK
dt = 1. / 240
link_id = robot.get_end_effector_ids(end_effector=0)
joint_ids = robot.joints # actuated joint
# joint_ids = joint_ids[2:]
damping = 0.01 # for damped-least-squares IK
wrt_link_id = -1 # robot.get_link_ids('iiwa_link_1')
# desired position
xd = np.array([0.5, 0., 0.5])
world.load_visual_sphere(xd, radius=0.05, color=(1, 0, 0, 0.5))
# change the robot visual
robot.change_transparency()
robot.draw_link_frames(link_ids=[-1, 0])
robot.draw_bounding_boxes(link_ids=joint_ids[0])
# robot.draw_link_coms([-1,0])
qIdx = robot.get_q_indices(joint_ids)
# OPTION 1: using `calculate_inverse_kinematics`###
if solver_flag == 0:
for _ in count():
# get current position in the task/operational space
x = robot.get_link_world_positions(link_id)
# print("(xd - x) = {}".format(xd - x))
last_trigger = controller.last_updated_button
movement = controller[last_trigger]
# print("Last updated button: {} with value: {}".format(last_trigger, movement))
if last_trigger == 'LJ':
X_desired[0] = X_desired[0] + round(movement[0],
1) * speed_scale_factor
X_desired[1] = X_desired[1] + round(movement[1],
1) * speed_scale_factor
elif last_trigger == 'RJ':
X_desired[2] = X_desired[2] + round(movement[1],
1) * speed_scale_factor
if sphere:
world.sim.remove_body(sphere)
sphere = world.load_visual_sphere(X_desired,
radius=0.005,
color=(1, 0, 0, 0.5))
## Move EE
# get current position in the task/operational space
# x = robot.get_link_world_positions(link_id)
# x = robot.sim.get_link_world_positions(body_id=robot.id, link_ids=link_id)
# # print("(xd - x) = {}".format(xd - x))
#
# # perform full IK
# q = robot.calculate_inverse_kinematics(link_id, position=X_desired)
#
# # set the joint positions
# robot.set_joint_positions(q[qIdx], joint_ids)
#
# # step in simulation
class DemoRecEnv:
def __init__(self, demo_dt, demo_dir):
# Initialize XBox controller
self.controller = XboxControllerInterface(use_thread=True,
verbose=False)
# Create simulator
self.sim = Bullet()
# create world
self.world = BasicWorld(self.sim)
# Create Robot
self.robot = Robot(self.sim, self.world)
self.robot.go_home()
self.workspace_outer_bound = {'rad': 1.0}
self.workspace_inner_bound = {'rad': 0.25, 'height': 0.7}
self.env_objects = []
self.demo_dir = demo_dir
self.demo_dt = demo_dt
def init_task_recorder(self):
num_viapts = 2 # int(input("How many via points?"))
via_pts = []
for i in range(num_viapts):
pt = generate_pt(via_pts, self.workspace_outer_bound,
self.workspace_inner_bound)
via_pts.append(pt)
self.env_objects.append(
self.world.load_visual_sphere(pt,
radius=0.05,
color=(1, 0, 0, 1)))
print("Via-pts generated.")
print(
"Press 'start' to start/stop recording demonstrations,select to end training and Y to reset via pts"
)
while True:
last_trigger = self.controller.last_updated_button
trigger_value = self.controller[last_trigger]
if last_trigger == 'menu' and trigger_value == 1:
self.start_recording()
elif last_trigger == 'view' and trigger_value == 1:
print("Ending Training")
self.reset_env()
break
elif last_trigger == 'Y' and trigger_value == 1:
for obj in self.env_objects:
self.world.sim.remove_body(obj)
self.init_task_recorder()
return
def start_recording(self):
time.sleep(1)
rec_traj = []
print(
"Started Recording. Press 'start' again to stop recording and 'A' to send robot to end"
)
sphere = None
X_desired = self.robot.robot.sim.get_link_world_positions(
body_id=self.robot.robot.id, link_ids=self.robot.ee_id)
speed_scale_factor = 0.01
time_ind = 0
save_every = 3
for t in count():
last_trigger = self.controller.last_updated_button
trigger_value = self.controller[last_trigger]
movement = self.controller[last_trigger]
if last_trigger == 'menu' and trigger_value == 1:
print("Stopping recording and resetting environment")
self.save_traj(rec_traj)
self.reset_env()
elif last_trigger == 'A' and trigger_value == 1:
X_desired = self.robot.end_pos
# Update target robot location from controller input
elif last_trigger == 'LJ':
X_desired[0] = X_desired[0] + round(movement[0],
1) * speed_scale_factor
X_desired[1] = X_desired[1] + round(movement[1],
1) * speed_scale_factor
elif last_trigger == 'RJ':
X_desired[2] = X_desired[2] + round(movement[1],
1) * speed_scale_factor
# Record points every 10 iterations
if not t % save_every:
x = self.robot.robot.sim.get_link_world_positions(
body_id=self.robot.robot.id, link_ids=self.robot.ee_id)
dx = self.robot.robot.get_link_world_linear_velocities(
self.robot.ee_id)
rec_traj.append((time_ind * self.demo_dt, ) +
tuple(round(i, 3) for i in x) +
tuple(round(i, 3) for i in dx))
time_ind = time_ind + 1
# Visualize target EE location
if sphere:
self.world.sim.remove_body(sphere)
sphere = self.world.load_visual_sphere(X_desired,
radius=0.01,
color=(0, 1, 0, 1))
# Send robot to target location
self.robot.send_robot_to(X_desired)
self.world.step()
def reset_env(self):
self.robot.go_home()
self.remove_objects()
self.init_task_recorder()
def remove_objects(self):
for obj in self.env_objects:
self.world.sim.remove_body(obj)
def step_env(self):
self.controller.step()
self.world.step()
def save_traj(self, rec_traj):
# Finding the index of last file saved
print(rec_traj[:10])
listdir = os.listdir(self.demo_dir)
listdir = list(map(int, listdir))
listdir.sort()
# print(listdir)
ind_last_file = 0 if len(listdir) == 0 else int(listdir[-1])
# print(ind_last_file)
del (rec_traj[0])
with open(self.demo_dir + str(ind_last_file + 1), 'w') as f:
write = csv.writer(f)
write.writerows(rec_traj)
class Environment:
def __init__(self, spheres=[[0.09, -0.44, 0.715, 1]], cuboids=None):
# Define size of the map
self.res = 1.
# Initialize XBox controller
self.controller = XboxControllerInterface(use_thread=True,
verbose=False)
# Create simulator
self.sim = Bullet()
# create world
self.world = BasicWorld(self.sim)
# create robot
self.robot = KukaIIWA(self.sim)
# define start/end pt
self.start_pt = [-0.75, 0., 0.75]
self.end_pt = [0.75, 0., 0.75]
# Initialize robot location
self.send_robot_to(np.array(self.start_pt), dt=2)
# Get via points from user
via_pts = self.get_via_pts(
) # [[-0.19, -0.65, 0.77], [0.51, -0.44, 0.58]]
self.motion_targets = [self.start_pt] + via_pts + [self.end_pt]
def send_robot_to(self, location, dt):
start = time.perf_counter()
for _ in count():
joint_ids = self.robot.joints
link_id = self.robot.get_end_effector_ids(end_effector=0)
qIdx = self.robot.get_q_indices(joint_ids)
x = self.robot.sim.get_link_world_positions(body_id=self.robot.id,
link_ids=link_id)
q = self.robot.calculate_inverse_kinematics(link_id,
position=location)
self.robot.set_joint_positions(q[qIdx], joint_ids)
self.world.step()
if time.perf_counter() - start >= dt:
break
def get_via_pts(self):
print(
"Move green dot to via point. Press 'A' to save. Press 'X' to finish."
)
X_desired = [0., 0., 1.3]
speed_scale_factor = 0.01
via_pts = []
sphere = None
for _ in count():
last_trigger = self.controller.last_updated_button
movement = self.controller[last_trigger]
if last_trigger == 'LJ':
X_desired[0] = X_desired[0] + round(movement[0],
1) * speed_scale_factor
X_desired[1] = X_desired[1] + round(movement[1],
1) * speed_scale_factor
elif last_trigger == 'RJ':
X_desired[2] = X_desired[2] + round(movement[1],
1) * speed_scale_factor
elif last_trigger == 'A' and movement == 1:
print("Via point added")
via_pts.append(X_desired)
time.sleep(0.5)
elif last_trigger == 'X' and movement == 1:
self.world.sim.remove_body(sphere)
print("Via point generation complete.")
break
if sphere:
self.world.sim.remove_body(sphere)
sphere = self.world.load_visual_sphere(X_desired,
radius=0.01,
color=(0, 1, 0, 1))
return via_pts
