def make_robosuite(env_name, env_config):
import robosuite
from robosuite.wrappers import IKWrapper
env = robosuite.make(
env_name,
has_renderer=env_config.render,
ignore_done=True,
use_camera_obs=env_config.pixel_input,
has_offscreen_renderer=env_config.pixel_input,
camera_height=84,
camera_width=84,
render_collision_mesh=False,
render_visual_mesh=True,
camera_name='agentview',
use_object_obs=(not env_config.pixel_input),
camera_depth=env_config.use_depth,
reward_shaping=True,
# demo_config=env_config.demonstration,
)
env = IKWrapper(env)
env = RobosuiteWrapper(env, env_config)
env = FilterWrapper(env, env_config)
env = ObservationConcatenationWrapper(env)
if env_config.pixel_input:
env = TransposeWrapper(env)
if env_config.use_grayscale:
env = GrayscaleWrapper(env)
if env_config.frame_stacks:
env = FrameStackWrapper(env, env_config)
env_config.action_spec = env.action_spec()
env_config.obs_spec = env.observation_spec()
return env, env_config
def __init__(self, has_display=False):
# Create environment.
print("Do I have a display?", has_display)
# self.base_env = robosuite.make('BaxterLift', has_renderer=has_display)
self.base_env = robosuite.make("BaxterViz", has_renderer=has_display)
# Create kinematics object.
self.baxter_IK_object = IKWrapper(self.base_env)
self.environment = self.baxter_IK_object.env
def __init__(self, has_display=False):
# Create environment.
print("Do I have a display?", has_display)
# self.base_env = robosuite.make('BaxterLift', has_renderer=has_display)
# self.base_env = robosuite.make("SawyerViz", has_renderer=has_display)
self.base_env = SawyerViz(has_renderer=has_display)
# Create kinematics object.
# TODO (chongyi zheng): Do we need this?
self.sawyer_IK_object = IKWrapper(self.base_env)
self.environment = self.sawyer_IK_object.env
W = 640
CAM_ID = 2 # 0 front view, 1 bird view, 2 agent view
env = SawyerGrasp(
has_renderer=True, # no on-screen renderer
has_offscreen_renderer=
True, # off-screen renderer is required for camera observations
ignore_done=True, # (optional) never terminates episode
use_camera_obs=True, # use camera observations
camera_height=H, # set camera height
camera_width=W, # set camera width
camera_name='agentview', # use "agentview" camera
use_object_obs=True, # no object feature when training on pixels
camera_depth=True,
target_object='cereal')
env = IKWrapper(env)
# reset the environment
env.reset()
f, cx, cy = env._get_cam_K(CAM_ID)
K = np.array([[f, 0, cx], [0, f, cy], [0, 0, 1]])
print(K)
# Get original camera location
cam_pos, cam_quat = env._get_cam_pos(CAM_ID)
cam_rot = T.quat2mat(cam_quat)
cam_pose = np.eye(4)
cam_pose[:3, :3] = cam_rot
cam_pose[:3, 3] = cam_pos
print(cam_pose)
)
pre_env = robosuite.make(
args.environment,
ignore_done=True,
use_camera_obs=False,
has_renderer=True,
control_freq=100,
gripper_visualization=True,
box_pos=[0.63522776, -0.3287869, 0.82162434], # shift2
box_quat=[0.6775825618903728, 0, 0, 0.679425538604203], # shift2
)
# enable controlling the end effector directly instead of using joint velocities
pre_env = GymWrapper(pre_env)
env = IKWrapper(env) # note cannot disable this or things go wack
# wrap the environment with data collection wrapper
tmp_directory = "~/Robotics/{}".format(str(time.time()).replace(
".", "_")) # Change from temp to Robotics folder
env = DataCollectionWrapperBaseline(env, tmp_directory)
# initialize device
if args.device == "keyboard":
from robosuite.devices import Keyboard
device = Keyboard()
env.viewer.add_keypress_callback("any", device.on_press)
env.viewer.add_keyup_callback("any", device.on_release)
env.viewer.add_keyrepeat_callback("any", device.on_press)
elif args.device == "spacemouse":
default="SawyerPickPlaceCan")
parser.add_argument("--device", type=str, default="keyboard")
args = parser.parse_args()
env = robosuite.make(
args.environment,
has_renderer=True,
ignore_done=True,
use_camera_obs=False,
gripper_visualization=True,
reward_shaping=True,
control_freq=100,
)
# enable controlling the end effector directly instead of using joint velocities
env = IKWrapper(env)
# initialize device
if args.device == "keyboard":
from robosuite.devices import Keyboard
device = Keyboard()
env.viewer.add_keypress_callback("any", device.on_press)
env.viewer.add_keyup_callback("any", device.on_release)
env.viewer.add_keyrepeat_callback("any", device.on_press)
elif args.device == "spacemouse":
from robosuite.devices import SpaceMouse
device = SpaceMouse()
else:
raise Exception(
type=bool,
default=False,
)
args = parser.parse_args()
# initialize a Baxter environment
env = robosuite.make(
"BaxterLine",
ignore_done=True,
has_renderer=args.render,
gripper_visualization=True,
use_camera_obs=False,
)
# env.viewer.set_camera(camera_id=2)
env = IKWrapper(env, 1)
tmp_directory = "/tmp/{}".format(str(time.time()).replace(".", "_"))
env = DataCollectionWrapper(env, tmp_directory)
j_states = np.load('line_start_pos.npz')
for i, g in enumerate(j_states['goals']):
if i >= 3:
run_collection(env, j_states['goals'][i - 3][1],
j_states['jpos'][i], g, True)
run_collection(env, j_states['goals'][i - 3][1],
j_states['jpos'][i], g, False)
if i <= len(j_states['goals']) - 4:
run_collection(env, j_states['goals'][i + 3][1],
j_states['jpos'][i], g, True)
run_collection(env, j_states['goals'][i + 3][1],
placement_initializer=True,
has_renderer=preview_mode,
ignore_done=True,
use_camera_obs=(not preview_mode),
camera_name="birdview",
camera_height=64,
camera_width=64,
gripper_visualization=True,
reward_shaping=True,
control_freq=20,
controller="position",
camera_depth=True,
)
# IK controller: we only use this IK, not control
ik_controller = IKWrapper(env).controller
trajectories_file = fannypack.utils.TrajectoriesFile(target_path,
read_only=False)
while len(trajectories_file) < args.traj_count:
obs = env.reset()
if preview_mode:
env.render()
env.controller.step = 0.0
env.controller.last_goal_position = np.array((0, 0, 0))
env.controller.last_goal_orientation = np.eye(3)
# Initialize training policy
policy: waypoint_policies.AbstractWaypointPolicy
if policy_mode == "push":
parser.add_argument("--device", type=str, default="keyboard")
args = parser.parse_args()
env = robosuite.make(
args.environment,
has_renderer=True,
ignore_done=True,
use_camera_obs=False,
gripper_visualization=True,
reward_shaping=True,
control_freq=100,
interpolation=None,
)
# enable controlling the end effector directly instead of using joint velocities
env = IKWrapper(env)
controller = env.controller
print("Controller: {}".format(controller))
# initialize device
if args.device == "keyboard":
from robosuite.devices import Keyboard
device = Keyboard()
env.viewer.add_keypress_callback("any", device.on_press)
env.viewer.add_keyup_callback("any", device.on_release)
env.viewer.add_keyrepeat_callback("any", device.on_press)
elif args.device == "spacemouse":
from robosuite.devices import SpaceMouse
device = SpaceMouse()
import numpy as np
import robosuite
from robosuite.wrappers import IKWrapper
if __name__ == "__main__":
# initialize a Baxter environment
env = robosuite.make(
"BaxterLift",
ignore_done=True,
has_renderer=True,
gripper_visualization=True,
use_camera_obs=False,
)
env = IKWrapper(env)
obs = env.reset()
# rotate the gripper so we can see it easily
env.set_robot_joint_positions([
0.00,
-0.55,
0.00,
1.28,
0.00,
0.26,
0.00,
0.00,
-0.55,
0.00,
def teleoperate(self, demons_config):
env = self.get_env()
# Need to use inverse-kinematics controller to set position using device
env = IKWrapper(env)
if demons_config.device == "keyboard":
from robosuite.devices import Keyboard
device = Keyboard()
env.viewer.add_keypress_callback("any", device.on_press)
env.viewer.add_keyup_callback("any", device.on_release)
env.viewer.add_keyrepeat_callback("any", device.on_press)
elif demons_config.device == "spacemouse":
from robosuite.devices import SpaceMouse
device = SpaceMouse()
for run in range(demons_config.n_runs):
obs = env.reset()
env.set_robot_joint_positions([0, -1.18, 0.00, 2.18, 0.00, 0.57, 1.5708])
# rotate the gripper so we can see it easily - NOTE : REMOVE MAYBE
env.viewer.set_camera(camera_id=2)
env.render()
device.start_control()
reset = False
task_completion_hold_count = -1
step = 0
tr_vobvs, tr_dof, tr_actions = [], [], []
while not reset:
if int(step % demons_config.collect_freq) == 0:
tr_vobvs.append(np.array(obs[self.vis_obv_key]))
tr_dof.append(np.array(obs[self.dof_obv_key].flatten()))
device_state = device.get_controller_state()
dpos, rotation, grasp, reset = (
device_state["dpos"],
device_state["rotation"],
device_state["grasp"],
device_state["reset"],
)
current = env._right_hand_orn
drotation = current.T.dot(rotation)
dquat = T.mat2quat(drotation)
grasp = grasp - 1.
ik_action = np.concatenate([dpos, dquat, [grasp]])
obs, _, done, _ = env.step(ik_action)
env.render()
joint_velocities = np.array(env.controller.commanded_joint_velocities)
if env.env.mujoco_robot.name == "sawyer":
gripper_actuation = np.array(ik_action[7:])
elif env.env.mujoco_robot.name == "baxter":
gripper_actuation = np.array(ik_action[14:])
# NOTE: Action for the normal environment (not inverse kinematic)
action = np.concatenate([joint_velocities, gripper_actuation], axis=0)
if int(step % demons_config.collect_freq) == 0:
tr_actions.append(action)
if (int(step % demons_config.flush_freq) == 0) or (demons_config.break_traj_success and task_completion_hold_count == 0):
print("Storing Trajectory")
trajectory = {self.vis_obv_key : np.array(tr_vobvs), self.dof_obv_key : np.array(tr_dof), 'action' : np.array(tr_actions)}
store_trajectoy(trajectory, 'teleop')
trajectory, tr_vobvs, tr_dof, tr_actions = {}, [], [], []
if demons_config.break_traj_success and env._check_success():
if task_completion_hold_count > 0:
task_completion_hold_count -= 1 # latched state, decrement count
else:
task_completion_hold_count = 10 # reset count on first success timestep
else:
task_completion_hold_count = -1
step += 1
env.close()
class robosuite_IKmover():
def __init__(self, env):
self.env_init = env
self.env = IKWrapper(env)
self.grasp = -1
def move_orient(self, env, q1):
q0 = env.observation_spec()['eef_quat']
x_target = env.observation_spec()['eef_pos']
fraction_size = 50
for i in range(fraction_size):
q_target = T.quat_slerp(q0, q1, fraction=(i + 1) / fraction_size)
steps = 0
lamda = 0.01
current = env._right_hand_orn
target_rotation = T.quat2mat(q_target)
drotation = current.T.dot(target_rotation)
while (np.linalg.norm(T.mat2euler(drotation)) > 0.01):
current = env._right_hand_orn
target_rotation = T.quat2mat(q_target)
drotation = current.T.dot(target_rotation)
dquat = T.mat2quat(drotation)
x_current = env.observation_spec()['eef_pos']
d_pos = np.clip((x_target - x_current) * lamda, -0.05, 0.05)
d_pos = [0, 0, 0]
action = np.concatenate((d_pos, dquat, [self.grasp]))
env.step(action)
env.render()
steps += 1
if (steps > 20):
break
return
def move_xy(self, env, x_target, target_rotation):
# target_rotation = np.array([[-1., 0., 0.], [0., 1., 0.], [0., 0., -1.]])
x_current = env.observation_spec()['eef_pos'][:2]
steps = 0
lamda = 0.1
while (np.linalg.norm(x_target - x_current) > 0.00001):
if (np.linalg.norm(x_target - x_current) < 0.01):
lamda = lamda * 1
else:
lamda = 0.1
x_current = env.observation_spec()['eef_pos'][:2]
current = env._right_hand_orn
drotation = current.T.dot(target_rotation)
dquat = T.mat2quat(drotation)
d_pos = (x_target - x_current) * lamda
action = np.concatenate((d_pos, [0], dquat, [self.grasp]))
env.step(action)
env.render()
for i in range(4):
# Now do action for zero xyz change so that bot stabilizes
current = env._right_hand_orn
drotation = current.T.dot(target_rotation)
dquat = T.mat2quat(drotation)
action = np.concatenate(([0, 0, 0], dquat, [self.grasp]))
env.step(action)
env.render()
steps += 1
if (steps > 500):
break
return
def move_z(self, env, z_target, target_rotation):
# target_rotation = np.array([[-1., 0., 0.], [0., 1., 0.], [0., 0., -1.]])
x_current = env.observation_spec()['eef_pos']
x_target = np.copy(x_current)
x_target[2] = z_target
steps = 0
lamda = 0.1
while (np.linalg.norm(x_target - x_current) > 0.00001):
if (np.linalg.norm(x_target - x_current) < 0.01):
lamda = lamda * 1.01
else:
lamda = 0.1
x_current = env.observation_spec()['eef_pos']
current = env._right_hand_orn
drotation = current.T.dot(target_rotation)
dquat = T.mat2quat(drotation)
d_pos = (x_target - x_current) * lamda
action = np.concatenate((d_pos, dquat, [self.grasp]))
env.step(action)
env.render()
for i in range(4):
# Now do action for zero xyz change so that bot stabilizes
current = env._right_hand_orn
drotation = current.T.dot(target_rotation)
dquat = T.mat2quat(drotation)
action = np.concatenate(([0, 0, 0], dquat, [self.grasp]))
env.step(action)
env.render()
steps += 1
if (steps > 500):
break
return
def move_xyz(self, env, x_target, target_rotation):
# target_rotation = np.array([[-1., 0., 0.], [0., 1., 0.], [0., 0., -1.]])
x_current = env.observation_spec()['eef_pos']
steps = 0
lamda = 0.1
while (np.linalg.norm(x_target - x_current) > 0.00001):
if (np.linalg.norm(x_target - x_current) < 0.01):
lamda = lamda * 1.01
else:
lamda = 0.1
x_current = env.observation_spec()['eef_pos']
current = env._right_hand_orn
drotation = current.T.dot(target_rotation)
dquat = T.mat2quat(drotation)
d_pos = (x_target - x_current) * lamda
action = np.concatenate((d_pos, dquat, [self.grasp]))
env.step(action)
env.render()
for i in range(4):
# Now do action for zero xyz change so that bot stabilizes
current = env._right_hand_orn
drotation = current.T.dot(target_rotation)
dquat = T.mat2quat(drotation)
action = np.concatenate(([0, 0, 0], dquat, [self.grasp]))
env.step(action)
env.render()
steps += 1
if (steps > 500):
break
return
def move(self, pos, rotation):
# First align the gripper angle
q1 = T.mat2quat(rotation)
self.move_orient(self.env, q1)
# Now move to x-y coordinate given by the pos
self.move_xy(self.env, pos[:2], rotation)
# Now move to x-y-z target position
self.move_xyz(self.env, pos, rotation)
def lift(self, pos):
print("Lifting")
rotation = T.quat2mat(self.env.observation_spec()['eef_quat'])
self.grasp = 1
self.move_z(self.env, z_target=pos[2], target_rotation=rotation)
self.move_xyz(self.env, x_target=pos, target_rotation=rotation)
def __init__(self, env):
self.env_init = env
self.env = IKWrapper(env)
self.grasp = -1