def _reset_internal(self):
"""
Resets simulation internal configurations.
"""
super()._reset_internal()
# Reset all object positions using initializer sampler if we're not directly loading from an xml
if not self.deterministic_reset:
# Sample from the placement initializer for all objects
object_placements = self.placement_initializer.sample()
# Loop through all objects and reset their positions
for obj_pos, obj_quat, obj in object_placements.values():
# If prehensile, set the object normally
if self.prehensile:
self.sim.data.set_joint_qpos(
obj.joints[0],
np.concatenate([np.array(obj_pos),
np.array(obj_quat)]))
# Else, set the object in the hand of the robot and loop a few steps to guarantee the robot is grasping
# the object initially
else:
eef_rot_quat = T.mat2quat(
T.euler2mat(
[np.pi - T.mat2euler(self._eef0_xmat)[2], 0, 0]))
obj_quat = T.quat_multiply(obj_quat, eef_rot_quat)
for j in range(100):
# Set object in hand
self.sim.data.set_joint_qpos(
obj.joints[0],
np.concatenate(
[self._eef0_xpos,
np.array(obj_quat)]))
# Close gripper (action = 1) and prevent arm from moving
if self.env_configuration == 'bimanual':
# Execute no-op action with gravity compensation
torques = np.concatenate([
self.robots[0].controller["right"].
torque_compensation, self.robots[0].
controller["left"].torque_compensation
])
self.sim.data.ctrl[self.robots[
0]._ref_joint_actuator_indexes] = torques
# Execute gripper action
self.robots[0].grip_action(
gripper=self.robots[0].gripper["right"],
gripper_action=[1])
else:
# Execute no-op action with gravity compensation
self.sim.data.ctrl[self.robots[0]._ref_joint_actuator_indexes] =\
self.robots[0].controller.torque_compensation
self.sim.data.ctrl[self.robots[1]._ref_joint_actuator_indexes] = \
self.robots[1].controller.torque_compensation
# Execute gripper action
self.robots[0].grip_action(
gripper=self.robots[0].gripper,
gripper_action=[1])
# Take forward step
self.sim.step()
def get_controller_state(self):
"""Returns the current state of the device, a dictionary of pos, orn, grasp, and reset."""
# Read 6 DOF pose from the tracker
pose = self._tracker.get_pose_euler()
# The position and orientation needs a coordinate transformation due to conventions used
# by the HTC Vive System and SteamVR
position = np.array([-pose[2], -pose[0], pose[1]]) - self._origin
# Note orientations are reported in degrees
orientation_rpy = np.deg2rad(pose[3:])
orientation_rpy = orientation_rpy[[0, 2, 1]]
orientation_rpy[0] *= -1
orientation_rpy[1] = wrap_to_pi(orientation_rpy[1] - np.pi / 2)
orientation_rpy[1] *= -1
# orientation_rpy = orientation_rpy[[1, 0, 2]]
rotation = euler2mat(orientation_rpy)
# rotation = self._initial_ori
# orientation_rpy = mat2euler(rotation)[[2, 1, 0]]
return dict(
dpos=position,
rotation=rotation,
raw_drotation=orientation_rpy,
grasp=0, # We don't control the grasp using the Vive Tracker
reset=
False # We don't control the reset signal using the Vive Tracker
)
def set_base_ori(self, rot):
"""
Rotates robot by rotation @rot from its original orientation.
Args:
rot (3-array): (r,p,y) euler angles specifying the orientation for the robot base
"""
# xml quat assumes w,x,y,z so we need to convert to this format from outputted x,y,z,w format from fcn
rot = mat2quat(euler2mat(rot))[[3, 0, 1, 2]]
self._elements["root_body"].set("quat", array_to_string(rot))
def _load_model(self):
"""
Loads an xml model, puts it in self.model
"""
super()._load_model()
# Adjust base pose(s) accordingly
if self.env_configuration == "bimanual":
xpos = self.robots[0].robot_model.base_xpos_offset["table"](
self.table_full_size[0])
self.robots[0].robot_model.set_base_xpos(xpos)
else:
if self.env_configuration == "single-arm-opposed":
# Set up robots facing towards each other by rotating them from their default position
for robot, rotation in zip(self.robots,
(np.pi / 2, -np.pi / 2)):
xpos = robot.robot_model.base_xpos_offset["table"](
self.table_full_size[0])
rot = np.array((0, 0, rotation))
xpos = T.euler2mat(rot) @ np.array(xpos)
robot.robot_model.set_base_xpos(xpos)
robot.robot_model.set_base_ori(rot)
else: # "single-arm-parallel" configuration setting
# Set up robots parallel to each other but offset from the center
for robot, offset in zip(self.robots, (-0.25, 0.25)):
xpos = robot.robot_model.base_xpos_offset["table"](
self.table_full_size[0])
xpos = np.array(xpos) + np.array((0, offset, 0))
robot.robot_model.set_base_xpos(xpos)
# load model for table top workspace
self.mujoco_arena = TableArena(
table_full_size=self.table_full_size,
table_friction=self.table_friction,
table_offset=(0, 0, 0.8),
)
if self.use_indicator_object:
self.mujoco_arena.add_pos_indicator()
# Arena always gets set to zero origin
self.mujoco_arena.set_origin([0, 0, 0])
# initialize objects of interest
self.pot = PotWithHandlesObject(name="pot")
self.mujoco_objects = OrderedDict([("pot", self.pot)])
# task includes arena, robot, and objects of interest
self.model = ManipulationTask(
mujoco_arena=self.mujoco_arena,
mujoco_robots=[robot.robot_model for robot in self.robots],
mujoco_objects=self.mujoco_objects,
visual_objects=None,
initializer=self.placement_initializer,
)
self.model.place_objects()
def set_base_ori(self, rot):
"""
Rotates robot by rotation @rot from its original orientation.
Args:
rot (3-array): (r,p,y) euler angles specifying the orientation for the robot base
"""
node = self.worldbody.find("./body[@name='{}']".format(self._root_))
# xml quat assumes w,x,y,z so we need to convert to this format from outputted x,y,z,w format from fcn
rot = mat2quat(euler2mat(rot))[[3, 0, 1, 2]]
node.set("quat", array_to_string(rot))
def get_interpolated_goal(self):
"""
Provides the next step in interpolation given the remaining steps.
NOTE: If this interpolator is for orientation, it is assumed to be receiving either euler angles or quaternions
Returns:
np.array: Next position in the interpolated trajectory
"""
# Grab start position
x = np.array(self.start)
# Calculate the desired next step based on remaining interpolation steps
if self.ori_interpolate is not None:
# This is an orientation interpolation, so we interpolate linearly around a sphere instead
goal = np.array(self.goal)
if self.ori_interpolate == "euler":
# this is assumed to be euler angles (x,y,z), so we need to first map to quat
x = T.mat2quat(T.euler2mat(x))
goal = T.mat2quat(T.euler2mat(self.goal))
# Interpolate to the next sequence
x_current = T.quat_slerp(x,
goal,
fraction=(self.step + 1) /
self.total_steps)
if self.ori_interpolate == "euler":
# Map back to euler
x_current = T.mat2euler(T.quat2mat(x_current))
else:
# This is a normal interpolation
dx = (self.goal - x) / (self.total_steps - self.step)
x_current = x + dx
# Increment step if there's still steps remaining based on ramp ratio
if self.step < self.total_steps - 1:
self.step += 1
# Return the new interpolated step
return x_current
def _load_model(self):
"""
Loads an xml model, puts it in self.model
"""
super()._load_model()
# Adjust base pose(s) accordingly
if self.env_configuration == "bimanual":
xpos = self.robots[0].robot_model.base_xpos_offset["table"](self.table_full_size[0])
self.robots[0].robot_model.set_base_xpos(xpos)
else:
if self.env_configuration == "single-arm-opposed":
# Set up robots facing towards each other by rotating them from their default position
for robot, rotation in zip(self.robots, (np.pi/2, -np.pi/2)):
xpos = robot.robot_model.base_xpos_offset["table"](self.table_full_size[0])
rot = np.array((0, 0, rotation))
xpos = T.euler2mat(rot) @ np.array(xpos)
robot.robot_model.set_base_xpos(xpos)
robot.robot_model.set_base_ori(rot)
else: # "single-arm-parallel" configuration setting
# Set up robots parallel to each other but offset from the center
for robot, offset in zip(self.robots, (-0.25, 0.25)):
xpos = robot.robot_model.base_xpos_offset["table"](self.table_full_size[0])
xpos = np.array(xpos) + np.array((0, offset, 0))
robot.robot_model.set_base_xpos(xpos)
# load model for table top workspace
mujoco_arena = TableArena(
table_full_size=self.table_full_size,
table_friction=self.table_friction,
table_offset=self.table_offset,
)
# Arena always gets set to zero origin
mujoco_arena.set_origin([0, 0, 0])
# initialize objects of interest
self.pot = PotWithHandlesObject(name="pot")
# Create placement initializer
if self.placement_initializer is not None:
self.placement_initializer.reset()
self.placement_initializer.add_objects(self.pot)
else:
self.placement_initializer = UniformRandomSampler(
name="ObjectSampler",
mujoco_objects=self.pot,
x_range=[-0.03, 0.03],
y_range=[-0.03, 0.03],
ensure_object_boundary_in_range=False,
ensure_valid_placement=True,
reference_pos=self.table_offset,
rotation=(np.pi + -np.pi / 3, np.pi + np.pi / 3),
)
# task includes arena, robot, and objects of interest
self.model = ManipulationTask(
mujoco_arena=mujoco_arena,
mujoco_robots=[robot.robot_model for robot in self.robots],
mujoco_objects=self.pot,
)
def _load_model(self):
"""
Loads an xml model, puts it in self.model
"""
super()._load_model()
# Adjust base pose(s) accordingly
if self.env_configuration == "bimanual":
xpos = self.robots[0].robot_model.base_xpos_offset["table"](
self.table_full_size[0])
self.robots[0].robot_model.set_base_xpos(xpos)
else:
if self.env_configuration == "single-arm-opposed":
# Set up robots facing towards each other by rotating them from their default position
for robot, rotation, offset in zip(self.robots,
(np.pi / 2, -np.pi / 2),
(-0.25, 0.25)):
xpos = robot.robot_model.base_xpos_offset["table"](
self.table_full_size[0])
rot = np.array((0, 0, rotation))
xpos = T.euler2mat(rot) @ np.array(xpos)
xpos += np.array((0, offset, 0))
robot.robot_model.set_base_xpos(xpos)
robot.robot_model.set_base_ori(rot)
else: # "single-arm-parallel" configuration setting
# Set up robots parallel to each other but offset from the center
for robot, offset in zip(self.robots, (-0.6, 0.6)):
xpos = robot.robot_model.base_xpos_offset["table"](
self.table_full_size[0])
xpos = np.array(xpos) + np.array((0, offset, 0))
robot.robot_model.set_base_xpos(xpos)
# load model for table top workspace
mujoco_arena = TableArena(table_full_size=self.table_true_size,
table_friction=self.table_friction,
table_offset=self.table_offset)
# Arena always gets set to zero origin
mujoco_arena.set_origin([0, 0, 0])
# Modify default agentview camera
mujoco_arena.set_camera(camera_name="agentview",
pos=[
0.8894354364730311, -3.481824231498976e-08,
1.7383813133506494
],
quat=[
0.6530981063842773, 0.2710406184196472,
0.27104079723358154, 0.6530979871749878
])
# initialize objects of interest
self.hammer = HammerObject(name="hammer")
# Create placement initializer
if self.placement_initializer is not None:
self.placement_initializer.reset()
self.placement_initializer.add_objects(self.hammer)
else:
# Set rotation about y-axis if hammer starts on table else rotate about z if it starts in gripper
rotation_axis = 'y' if self.prehensile else 'z'
self.placement_initializer = UniformRandomSampler(
name="ObjectSampler",
mujoco_objects=self.hammer,
x_range=[-0.1, 0.1],
y_range=[-0.05, 0.05],
rotation=None,
rotation_axis=rotation_axis,
ensure_object_boundary_in_range=False,
ensure_valid_placement=True,
reference_pos=self.table_offset,
)
# task includes arena, robot, and objects of interest
self.model = ManipulationTask(
mujoco_arena=mujoco_arena,
mujoco_robots=[robot.robot_model for robot in self.robots],
mujoco_objects=self.hammer,
)
def set_goal_orientation(delta,
current_orientation,
orientation_limit=None,
set_ori=None):
"""
Calculates and returns the desired goal orientation, clipping the result accordingly to @orientation_limits.
@delta and @current_orientation must be specified if a relative goal is requested, else @set_ori must be
an orientation matrix specified to define a global orientation
Args:
delta (np.array): Desired relative change in orientation, in axis-angle form [ax, ay, az]
current_orientation (np.array): Current orientation, in rotation matrix form
orientation_limit (None or np.array): 2d array defining the (min, max) limits of permissible orientation goal commands
set_ori (None or np.array): If set, will ignore @delta and set the goal orientation to this value
Returns:
np.array: calculated goal orientation in absolute coordinates
Raises:
ValueError: [Invalid orientation_limit shape]
"""
# directly set orientation
if set_ori is not None:
goal_orientation = set_ori
# otherwise use delta to set goal orientation
else:
# convert axis-angle value to rotation matrix
quat_error = trans.axisangle2quat(delta)
rotation_mat_error = trans.quat2mat(quat_error)
goal_orientation = np.dot(rotation_mat_error, current_orientation)
# check for orientation limits
if np.array(orientation_limit).any():
if orientation_limit.shape != (2, 3):
raise ValueError("Orientation limit should be shaped (2,3) "
"but is instead: {}".format(
orientation_limit.shape))
# Convert to euler angles for clipping
euler = trans.mat2euler(goal_orientation)
# Clip euler angles according to specified limits
limited = False
for idx in range(3):
if orientation_limit[0][idx] < orientation_limit[1][
idx]: # Normal angle sector meaning
if orientation_limit[0][idx] < euler[idx] < orientation_limit[
1][idx]:
continue
else:
limited = True
dist_to_lower = euler[idx] - orientation_limit[0][idx]
if dist_to_lower > np.pi:
dist_to_lower -= 2 * np.pi
elif dist_to_lower < -np.pi:
dist_to_lower += 2 * np.pi
dist_to_higher = euler[idx] - orientation_limit[1][idx]
if dist_to_lower > np.pi:
dist_to_higher -= 2 * np.pi
elif dist_to_lower < -np.pi:
dist_to_higher += 2 * np.pi
if dist_to_lower < dist_to_higher:
euler[idx] = orientation_limit[0][idx]
else:
euler[idx] = orientation_limit[1][idx]
else: # Inverted angle sector meaning
if orientation_limit[0][idx] < euler[idx] or euler[
idx] < orientation_limit[1][idx]:
continue
else:
limited = True
dist_to_lower = euler[idx] - orientation_limit[0][idx]
if dist_to_lower > np.pi:
dist_to_lower -= 2 * np.pi
elif dist_to_lower < -np.pi:
dist_to_lower += 2 * np.pi
dist_to_higher = euler[idx] - orientation_limit[1][idx]
if dist_to_lower > np.pi:
dist_to_higher -= 2 * np.pi
elif dist_to_lower < -np.pi:
dist_to_higher += 2 * np.pi
if dist_to_lower < dist_to_higher:
euler[idx] = orientation_limit[0][idx]
else:
euler[idx] = orientation_limit[1][idx]
if limited:
goal_orientation = trans.euler2mat(
np.array([euler[0], euler[1], euler[2]]))
return goal_orientation
def _load_model(self):
"""
Loads an xml model, puts it in self.model
"""
super()._load_model()
# Adjust base pose(s) accordingly
if self.env_configuration == "bimanual":
xpos = self.robots[0].robot_model.base_xpos_offset["empty"]
self.robots[0].robot_model.set_base_xpos(xpos)
else:
if self.env_configuration == "single-arm-opposed":
# Set up robots facing towards each other by rotating them from their default position
for robot, rotation in zip(self.robots,
(np.pi / 2, -np.pi / 2)):
xpos = robot.robot_model.base_xpos_offset["empty"]
rot = np.array((0, 0, rotation))
xpos = T.euler2mat(rot) @ np.array(xpos)
robot.robot_model.set_base_xpos(xpos)
robot.robot_model.set_base_ori(rot)
else: # "single-arm-parallel" configuration setting
# Set up robots parallel to each other but offset from the center
for robot, offset in zip(self.robots, (-0.25, 0.25)):
xpos = robot.robot_model.base_xpos_offset["empty"]
xpos = np.array(xpos) + np.array((0, offset, 0))
robot.robot_model.set_base_xpos(xpos)
# Add arena and robot
self.model = MujocoWorldBase()
self.mujoco_arena = EmptyArena()
if self.use_indicator_object:
self.mujoco_arena.add_pos_indicator()
self.model.merge(self.mujoco_arena)
for robot in self.robots:
self.model.merge(robot.robot_model)
# initialize objects of interest
self.hole = PlateWithHoleObject(name="hole", )
tex_attrib = {
"type": "cube",
}
mat_attrib = {
"texrepeat": "1 1",
"specular": "0.4",
"shininess": "0.1",
}
greenwood = CustomMaterial(
texture="WoodGreen",
tex_name="greenwood",
mat_name="greenwood_mat",
tex_attrib=tex_attrib,
mat_attrib=mat_attrib,
)
self.peg = CylinderObject(
name="peg",
size_min=(self.peg_radius[0], self.peg_length),
size_max=(self.peg_radius[1], self.peg_length),
material=greenwood,
rgba=[0, 1, 0, 1],
)
# Load hole object
self.hole_obj = self.hole.get_collision(site=True)
self.hole_obj.set("quat", "0 0 0.707 0.707")
self.hole_obj.set("pos", "0.11 0 0.17")
self.model.merge_asset(self.hole)
# Load peg object
self.peg_obj = self.peg.get_collision(site=True)
self.peg_obj.set("pos", array_to_string((0, 0, self.peg_length)))
self.model.merge_asset(self.peg)
# Depending on env configuration, append appropriate objects to arms
if self.env_configuration == "bimanual":
self.model.worldbody.find(".//body[@name='{}']".format(
self.robots[0].robot_model.eef_name["left"])).append(
self.hole_obj)
self.model.worldbody.find(".//body[@name='{}']".format(
self.robots[0].robot_model.eef_name["right"])).append(
self.peg_obj)
else:
self.model.worldbody.find(".//body[@name='{}']".format(
self.robots[1].robot_model.eef_name)).append(self.hole_obj)
self.model.worldbody.find(".//body[@name='{}']".format(
self.robots[0].robot_model.eef_name)).append(self.peg_obj)
def input2action(device, robot, active_arm="right", env_configuration=None):
"""
Converts an input from an active device into a valid action sequence that can be fed into an env.step() call
If a reset is triggered from the device, immediately returns None. Else, returns the appropriate action
Args:
device (Device): A device from which user inputs can be converted into actions. Can be either a Spacemouse or
Keyboard device class
robot (Robot): Which robot we're controlling
active_arm (str): Only applicable for multi-armed setups (e.g.: multi-arm environments or bimanual robots).
Allows inputs to be converted correctly if the control type (e.g.: IK) is dependent on arm choice.
Choices are {right, left}
env_configuration (str or None): Only applicable for multi-armed environments. Allows inputs to be converted
correctly if the control type (e.g.: IK) is dependent on the environment setup. Options are:
{bimanual, single-arm-parallel, single-arm-opposed}
Returns:
2-tuple:
- (None or np.array): Action interpreted from @device including any gripper action(s). None if we get a
reset signal from the device
- (None or int): 1 if desired close, -1 if desired open gripper state. None if get a reset signal from the
device
"""
state = device.get_controller_state()
# Note: Devices output rotation with x and z flipped to account for robots starting with gripper facing down
# Also note that the outputted rotation is an absolute rotation, while outputted dpos is delta pos
# Raw delta rotations from neutral user input is captured in raw_drotation (roll, pitch, yaw)
dpos, rotation, raw_drotation, grasp, reset = (
state["dpos"],
state["rotation"],
state["raw_drotation"],
state["grasp"],
state["reset"],
)
# If we're resetting, immediately return None
if reset:
return None, None
# Get controller reference
controller = robot.controller if not isinstance(
robot, Bimanual) else robot.controller[active_arm]
gripper_dof = robot.gripper.dof if not isinstance(
robot, Bimanual) else robot.gripper[active_arm].dof
# First process the raw drotation
drotation = raw_drotation[[1, 0, 2]]
if controller.name == 'IK_POSE':
# If this is panda, want to swap x and y axis
if isinstance(robot.robot_model, Panda):
drotation = drotation[[1, 0, 2]]
else:
# Flip x
drotation[0] = -drotation[0]
# Scale rotation for teleoperation (tuned for IK)
drotation *= 10
dpos *= 5
# relative rotation of desired from current eef orientation
# map to quat
drotation = T.mat2quat(T.euler2mat(drotation))
# If we're using a non-forward facing configuration, need to adjust relative position / orientation
if env_configuration == "single-arm-opposed":
# Swap x and y for pos and flip x,y signs for ori
dpos = dpos[[1, 0, 2]]
drotation[0] = -drotation[0]
drotation[1] = -drotation[1]
if active_arm == "left":
# x pos needs to be flipped
dpos[0] = -dpos[0]
else:
# y pos needs to be flipped
dpos[1] = -dpos[1]
# Lastly, map to axis angle form
drotation = T.quat2axisangle(drotation)
elif controller.name == 'OSC_POSE':
# Flip z
drotation[2] = -drotation[2]
# Scale rotation for teleoperation (tuned for OSC) -- gains tuned for each device
drotation = drotation * 1.5 if isinstance(device,
Keyboard) else drotation * 50
dpos = dpos * 75 if isinstance(device, Keyboard) else dpos * 125
else:
# No other controllers currently supported
print(
"Error: Unsupported controller specified -- Robot must have either an IK or OSC-based controller!"
)
# map 0 to -1 (open) and map 1 to 1 (closed)
grasp = 1 if grasp else -1
# Create action based on action space of individual robot
action = np.concatenate([dpos, drotation, [grasp] * gripper_dof])
# Return the action and grasp
return action, grasp
def _load_model(self):
"""
Loads an xml model, puts it in self.model
"""
super()._load_model()
# Adjust base pose(s) accordingly
if self.env_configuration == "bimanual":
xpos = self.robots[0].robot_model.base_xpos_offset["empty"]
self.robots[0].robot_model.set_base_xpos(xpos)
else:
if self.env_configuration == "single-arm-opposed":
# Set up robots facing towards each other by rotating them from their default position
for robot, rotation in zip(self.robots,
(np.pi / 2, -np.pi / 2)):
xpos = robot.robot_model.base_xpos_offset["empty"]
rot = np.array((0, 0, rotation))
xpos = T.euler2mat(rot) @ np.array(xpos)
robot.robot_model.set_base_xpos(xpos)
robot.robot_model.set_base_ori(rot)
else: # "single-arm-parallel" configuration setting
# Set up robots parallel to each other but offset from the center
for robot, offset in zip(self.robots, (-0.25, 0.25)):
xpos = robot.robot_model.base_xpos_offset["empty"]
xpos = np.array(xpos) + np.array((0, offset, 0))
robot.robot_model.set_base_xpos(xpos)
# Add arena and robot
mujoco_arena = EmptyArena()
# Arena always gets set to zero origin
mujoco_arena.set_origin([0, 0, 0])
# Modify default agentview camera
mujoco_arena.set_camera(camera_name="agentview",
pos=[
1.0666432116509934, 1.4903257668114777e-08,
2.0563394967349096
],
quat=[
0.6530979871749878, 0.27104058861732483,
0.27104055881500244, 0.6530978679656982
])
# initialize objects of interest
self.hole = PlateWithHoleObject(name="hole")
tex_attrib = {
"type": "cube",
}
mat_attrib = {
"texrepeat": "1 1",
"specular": "0.4",
"shininess": "0.1",
}
greenwood = CustomMaterial(
texture="WoodGreen",
tex_name="greenwood",
mat_name="greenwood_mat",
tex_attrib=tex_attrib,
mat_attrib=mat_attrib,
)
self.peg = CylinderObject(
name="peg",
size_min=(self.peg_radius[0], self.peg_length),
size_max=(self.peg_radius[1], self.peg_length),
material=greenwood,
rgba=[0, 1, 0, 1],
joints=None,
)
# Load hole object
hole_obj = self.hole.get_obj()
hole_obj.set("quat", "0 0 0.707 0.707")
hole_obj.set("pos", "0.11 0 0.17")
# Load peg object
peg_obj = self.peg.get_obj()
peg_obj.set("pos", array_to_string((0, 0, self.peg_length)))
# Append appropriate objects to arms
if self.env_configuration == "bimanual":
r_eef, l_eef = [
self.robots[0].robot_model.eef_name[arm]
for arm in self.robots[0].arms
]
r_model, l_model = [
self.robots[0].robot_model, self.robots[0].robot_model
]
else:
r_eef, l_eef = [
robot.robot_model.eef_name for robot in self.robots
]
r_model, l_model = [
self.robots[0].robot_model, self.robots[1].robot_model
]
r_body = find_elements(root=r_model.worldbody,
tags="body",
attribs={"name": r_eef},
return_first=True)
l_body = find_elements(root=l_model.worldbody,
tags="body",
attribs={"name": l_eef},
return_first=True)
r_body.append(peg_obj)
l_body.append(hole_obj)
# task includes arena, robot, and objects of interest
# We don't add peg and hole directly since they were already appended to the robots
self.model = ManipulationTask(
mujoco_arena=mujoco_arena,
mujoco_robots=[robot.robot_model for robot in self.robots],
)
# Make sure to add relevant assets from peg and hole objects
self.model.merge_assets(self.hole)
self.model.merge_assets(self.peg)
# To accommodate for multi-arm settings (e.g.: Baxter), we need to make sure to fill any extra action space
# Get total number of arms being controlled
n = 0
for robot in env.robots:
n += int(robot.action_dim / action_dim)
# Keep track of done variable to know when to break loop
count = 0
# Loop through controller space
while count < num_test_steps:
action = neutral.copy()
for i in range(steps_per_action):
if controller_name == 'EE_IK' and count > 2:
# Convert from euler angle to quat here since we're working with quats
angle = np.zeros(3)
angle[count - 3] = test_value
action[3:7] = T.mat2quat(T.euler2mat(angle))
else:
action[count] = test_value
total_action = np.tile(action, n)
env.step(total_action)
env.render()
for i in range(steps_per_rest):
total_action = np.tile(neutral, n)
env.step(total_action)
env.render()
count += 1
# Shut down this env before starting the next test
env.close()
def test_linear_interpolator():
for controller_name in ["EE_POS_ORI", "EE_IK"]:
for traj in ["pos", "ori"]:
# Define counter to increment timesteps for each trajectory
timesteps = [0, 0]
for interpolator in [None, "linear"]:
# Define numpy seed so we guarantee consistent starting pos / ori for each trajectory
np.random.seed(3)
# Define controller path to load
controller_path = os.path.join(
os.path.dirname(__file__), '../../robosuite',
'controllers/config/{}.json'.format(
controller_name.lower()))
with open(controller_path) as f:
controller_config = json.load(f)
controller_config["interpolation"] = interpolator
controller_config["ramp_ratio"] = 1.0
# Now, create a test env for testing the controller on
env = suite.make(
"Lift",
robots="Sawyer",
has_renderer=args.
render, # by default, don't use on-screen renderer for visual validation
has_offscreen_renderer=False,
use_camera_obs=False,
horizon=10000,
control_freq=20,
controller_configs=controller_config)
# Reset the environment
env.reset()
# Notify user a new trajectory is beginning
print(
"\nTesting controller {} with trajectory {} and interpolator={}..."
.format(controller_name, traj, interpolator))
# If rendering, set controller to front view to get best angle for viewing robot movements
if args.render:
env.viewer.set_camera(camera_id=0)
# Keep track of state of robot eef (pos, ori (euler))
initial_state = [
env.robots[0]._right_hand_pos,
T.mat2euler(env.robots[0]._right_hand_orn)
]
dstate = [
env.robots[0]._right_hand_pos - initial_state[0],
T.mat2euler(env.robots[0]._right_hand_orn) -
initial_state[1]
]
# Define the uniform trajectory action
if traj == "pos":
pos_act = pos_action_ik if controller_name.lower(
) == "ee_ik" else pos_action_osc
rot_act = T.mat2quat(T.euler2mat(
np.zeros(3))) if controller_name.lower(
) == "ee_ik" else np.zeros(3)
else:
pos_act = np.zeros(3)
rot_act = rot_action_ik if controller_name.lower(
) == "ee_ik" else rot_action_osc
# Compose the action
action = np.concatenate([pos_act, rot_act, [0]])
# Determine which trajectory we're executing
k = 0 if traj == "pos" else 1
j = 0 if not interpolator else 1
# Run trajectory until the threshold condition is met
while abs(dstate[k][indexes[k]]) < abs(thresholds[k]):
env.step(action)
if args.render:
env.render()
# Update timestep count and state
timesteps[j] += 1
dstate = [
env.robots[0]._right_hand_pos - initial_state[0],
T.mat2euler(env.robots[0]._right_hand_orn) -
initial_state[1]
]
# When finished, print out the timestep results
print("Completed trajectory. Took {} timesteps total.".format(
timesteps[j]))
# Shut down this env before starting the next test
env.close()
# Assert that the interpolated path is slower than the non-interpolated one
assert timesteps[1] > min_ratio * timesteps[0], "Error: Interpolated trajectory time should be longer " \
"than non-interpolated!"
# Tests completed!
print()
print("-" * 80)
print("All linear interpolator testing completed.\n")
import robosuite.utils.transform_utils as T
import argparse
# Define the threshold locations, delta values, and ratio #
# Translation trajectory
pos_y_threshold = 0.1
delta_pos_y = 0.01
pos_action_osc = [0, delta_pos_y * 40, 0]
pos_action_ik = [0, delta_pos_y, 0]
# Rotation trajectory
rot_r_threshold = np.pi / 2
delta_rot_r = 0.01
rot_action_osc = [-delta_rot_r * 40, 0, 0]
rot_action_ik = T.mat2quat(T.euler2mat([delta_rot_r * 5, 0, 0]))
# Concatenated thresholds and corresponding indexes (y = 1 in x,y,z; roll = 0 in r,p,y)
thresholds = [pos_y_threshold, rot_r_threshold]
indexes = [1, 0]
# Threshold ratio
min_ratio = 1.10
# Define arguments for this test
parser = argparse.ArgumentParser()
parser.add_argument("--render",
action="store_true",
help="Whether to render tests or run headless")
args = parser.parse_args()
def test_all_controllers():
for controller_name in controllers.keys():
# Define variables for each controller test
action_dim = controllers[controller_name][0]
num_test_steps = controllers[controller_name][1]
test_value = controllers[controller_name][2]
neutral = controllers[controller_name][3]
# Define controller path to load
controller_path = os.path.join(
os.path.dirname(__file__), '../../robosuite',
'controllers/config/{}.json'.format(controller_name.lower()))
with open(controller_path) as f:
controller_config = json.load(f)
# Now, create a test env for testing the controller on
env = suite.make(
"Lift",
robots="Sawyer",
has_renderer=args.
render, # use on-screen renderer for visual validation only if requested
has_offscreen_renderer=False,
use_camera_obs=False,
horizon=(steps_per_action + steps_per_rest) * num_test_steps,
controller_configs=controller_config)
print("Testing controller: {}...".format(controller_name))
env.reset()
# If rendering, set controller to front view to get best angle for viewing robot movements
if args.render:
env.viewer.set_camera(camera_id=0)
# get action range
action_min, action_max = env.action_spec
assert action_min.shape == action_max.shape
assert action_min.shape[0] == action_dim, "Expected {}, got {}".format(
action_dim, action_min.shape[0])
# Keep track of done variable to know when to break loop
done = False
count = 0
# Loop through controller space
while count < num_test_steps:
action = neutral.copy()
for i in range(steps_per_action):
if controller_name.lower() == 'ee_ik' and count > 2:
# Convert from euler angle to quat here since we're working with quats
angle = np.zeros(3)
angle[count - 3] = test_value
action[3:7] = T.mat2quat(T.euler2mat(angle))
else:
action[count] = test_value
env.step(action)
if args.render:
env.render()
for i in range(steps_per_rest):
env.step(neutral)
if args.render:
env.render()
count += 1
# Shut down this env before starting the next test
env.close()
# Tests passed!
print("All controller tests completed.")
