def __init__(self, xyz_dims=3, xyz_remap=[0, 1, 2], xyz_scale=[1, 1, 1]):
"""TODO: fill in other params"""
from robosuite.devices import SpaceMouse
self.xyz_dims = xyz_dims
self.xyz_remap = np.array(xyz_remap)
self.xyz_scale = np.array(xyz_scale)
self.device = SpaceMouse()
class SpaceMouseExpert(Expert):
def __init__(self, xyz_dims=3, xyz_remap=[0, 1, 2], xyz_scale=[1, 1, 1]):
"""TODO: fill in other params"""
from robosuite.devices import SpaceMouse
self.xyz_dims = xyz_dims
self.xyz_remap = np.array(xyz_remap)
self.xyz_scale = np.array(xyz_scale)
self.device = SpaceMouse()
def get_action(self, obs):
"""Must return (action, valid, reset, accept)"""
state = self.device.get_controller_state()
dpos, rotation, accept, reset = (
state["dpos"],
state["rotation"],
state["left_click"],
state["right_click"],
)
xyz = dpos[self.xyz_remap] * self.xyz_scale
a = xyz[:self.xyz_dims]
valid = not np.all(np.isclose(a, 0))
return (a, valid, reset, accept)
# Setup printing options for numbers
np.set_printoptions(formatter={"float": lambda x: "{0:0.3f}".format(x)})
# initialize device
if args.device == "keyboard":
from robosuite.devices import Keyboard
device = Keyboard(pos_sensitivity=args.pos_sensitivity, rot_sensitivity=args.rot_sensitivity)
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(pos_sensitivity=args.pos_sensitivity, rot_sensitivity=args.rot_sensitivity)
else:
raise Exception("Invalid device choice: choose either 'keyboard' or 'spacemouse'.")
while True:
# Reset the environment
obs = env.reset()
# Setup rendering
cam_id = 0
num_cam = len(env.sim.model.camera_names)
env.render()
# Initialize variables that should the maintained between resets
last_grasp = 0
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":
from robosuite.devices import SpaceMouse
device = SpaceMouse()
else:
raise Exception(
"Invalid device choice: choose either 'keyboard' or 'spacemouse'.")
# make a new timestamped directory
t1, t2 = str(time.time()).split(".")
new_dir = os.path.join(args.directory, "{}_{}".format(t1, t2))
os.makedirs(new_dir)
# Setup network
# ----------------------------------------
policy_hidden_size = 100 # matches default
def policy_fn(name, ob_space, ac_space, reuse=False):
return mlp_policy_sawyer.MlpPolicy(name=name,
from multiworld.envs.mujoco.sawyer_xyz.sawyer_hammer_6dof import SawyerHammer6DOFEnv
from multiworld.envs.mujoco.sawyer_xyz.sawyer_box_open_6dof import SawyerBoxOpen6DOFEnv
from multiworld.envs.mujoco.sawyer_xyz.sawyer_bin_picking_6dof import SawyerBinPicking6DOFEnv
from multiworld.envs.mujoco.sawyer_xyz.sawyer_peg_insertion_side_6dof import SawyerPegInsertionSide6DOFEnv
from multiworld.envs.mujoco.sawyer_xyz.sawyer_peg_insertion_topdown_6dof import SawyerPegInsertionTopdown6DOFEnv
from multiworld.envs.mujoco.sawyer_xyz.sawyer_peg_unplug_side_6dof import SawyerPegUnplugSide6DOFEnv
from robosuite.devices import SpaceMouse
from multiworld.envs.mujoco.utils import rotation
from robosuite.utils.transform_utils import mat2quat
from multiworld.envs.env_util import quat_to_zangle, zangle_to_quat
import gym
import multiworld
space_mouse = SpaceMouse()
env = SawyerPegUnplugSide6DOFEnv(random_init=True, obs_type='with_goal')
NDIM = env.action_space.low.size
lock_action = False
obs = env.reset()
action = np.zeros(10)
closed = False
while True:
done = False
env.render()
state = space_mouse.get_controller_state()
dpos, rotation, grasp, reset = (
state["dpos"],
state["rotation"],
"""
Should be run on a machine connected to a spacemouse
"""
import time
import Pyro4
from robosuite.devices import SpaceMouse
from rlkit.launchers import config
Pyro4.config.SERIALIZERS_ACCEPTED = set(
["pickle", "json", "marshal", "serpent"])
Pyro4.config.SERIALIZER = "pickle"
nameserver = Pyro4.locateNS(host=config.SPACEMOUSE_HOSTNAME)
uri = nameserver.lookup("example.greeting")
device_state = Pyro4.Proxy(uri)
device = SpaceMouse()
while True:
state = device.get_controller_state()
print(state)
time.sleep(0.1)
device_state.set_state(state)
# wrap the environment with data collection wrapper
tmp_directory = "/tmp/{}".format(str(time.time()).replace(".", "_"))
env = DataCollectionWrapper(env, tmp_directory)
# initialize devices
if args.device1 == "keyboard":
from robosuite.devices import Keyboard
device1 = Keyboard()
env.viewer.add_keypress_callback("any", device1.on_press)
env.viewer.add_keyup_callback("any", device1.on_release)
env.viewer.add_keyrepeat_callback("any", device1.on_press)
elif args.device1 == "spacemouse":
from robosuite.devices import SpaceMouse
device1 = SpaceMouse()
else:
raise Exception(
"Invalid device choice: choose either 'keyboard' or 'spacemouse'.")
if args.device2 == "keyboard":
from robosuite.devices import Keyboard
device2 = Keyboard()
env.viewer.add_keypress_callback("any", device2.on_press)
env.viewer.add_keyup_callback("any", device2.on_release)
env.viewer.add_keyrepeat_callback("any", device2.on_press)
elif args.device2 == "spacemouse":
from robosuite.devices import SpaceMouse
device2 = SpaceMouse()
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()
