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)
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"],
state["grasp"],
state["reset"],
)
# convert into a suitable end effector action for the environment
# current = env.get_mocap_quat()
# desired_quat = mat2quat(rotation)
# current_z = quat_to_zangle(current)
# desired_z = quat_to_zangle(desired_quat)
# # drotation = current.T.dot(rotation) # relative rotation of desired from current
"""
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)
