class TwoArmLift(TwoArmEnv):
"""
This class corresponds to the lifting task for two robot arms.
Args:
robots (str or list of str): Specification for specific robot arm(s) to be instantiated within this env
(e.g: "Sawyer" would generate one arm; ["Panda", "Panda", "Sawyer"] would generate three robot arms)
Note: Must be either 2 single single-arm robots or 1 bimanual robot!
env_configuration (str): Specifies how to position the robots within the environment. Can be either:
:`'bimanual'`: Only applicable for bimanual robot setups. Sets up the (single) bimanual robot on the -x
side of the table
:`'single-arm-parallel'`: Only applicable for multi single arm setups. Sets up the (two) single armed
robots next to each other on the -x side of the table
:`'single-arm-opposed'`: Only applicable for multi single arm setups. Sets up the (two) single armed
robots opposed from each others on the opposite +/-y sides of the table.
Note that "default" corresponds to either "bimanual" if a bimanual robot is used or "single-arm-opposed" if two
single-arm robots are used.
controller_configs (str or list of dict): If set, contains relevant controller parameters for creating a
custom controller. Else, uses the default controller for this specific task. Should either be single
dict if same controller is to be used for all robots or else it should be a list of the same length as
"robots" param
gripper_types (str or list of str): type of gripper, used to instantiate
gripper models from gripper factory. Default is "default", which is the default grippers(s) associated
with the robot(s) the 'robots' specification. None removes the gripper, and any other (valid) model
overrides the default gripper. Should either be single str if same gripper type is to be used for all
robots or else it should be a list of the same length as "robots" param
initialization_noise (dict or list of dict): Dict containing the initialization noise parameters.
The expected keys and corresponding value types are specified below:
:`'magnitude'`: The scale factor of uni-variate random noise applied to each of a robot's given initial
joint positions. Setting this value to `None` or 0.0 results in no noise being applied.
If "gaussian" type of noise is applied then this magnitude scales the standard deviation applied,
If "uniform" type of noise is applied then this magnitude sets the bounds of the sampling range
:`'type'`: Type of noise to apply. Can either specify "gaussian" or "uniform"
Should either be single dict if same noise value is to be used for all robots or else it should be a
list of the same length as "robots" param
:Note: Specifying "default" will automatically use the default noise settings.
Specifying None will automatically create the required dict with "magnitude" set to 0.0.
table_full_size (3-tuple): x, y, and z dimensions of the table.
table_friction (3-tuple): the three mujoco friction parameters for
the table.
use_camera_obs (bool): if True, every observation includes rendered image(s)
use_object_obs (bool): if True, include object (cube) information in
the observation.
reward_scale (None or float): Scales the normalized reward function by the amount specified.
If None, environment reward remains unnormalized
reward_shaping (bool): if True, use dense rewards.
placement_initializer (ObjectPositionSampler): if provided, will
be used to place objects on every reset, else a UniformRandomSampler
is used by default.
has_renderer (bool): If true, render the simulation state in
a viewer instead of headless mode.
has_offscreen_renderer (bool): True if using off-screen rendering
render_camera (str): Name of camera to render if `has_renderer` is True. Setting this value to 'None'
will result in the default angle being applied, which is useful as it can be dragged / panned by
the user using the mouse
render_collision_mesh (bool): True if rendering collision meshes in camera. False otherwise.
render_visual_mesh (bool): True if rendering visual meshes in camera. False otherwise.
render_gpu_device_id (int): corresponds to the GPU device id to use for offscreen rendering.
Defaults to -1, in which case the device will be inferred from environment variables
(GPUS or CUDA_VISIBLE_DEVICES).
control_freq (float): how many control signals to receive in every second. This sets the amount of
simulation time that passes between every action input.
horizon (int): Every episode lasts for exactly @horizon timesteps.
ignore_done (bool): True if never terminating the environment (ignore @horizon).
hard_reset (bool): If True, re-loads model, sim, and render object upon a reset call, else,
only calls sim.reset and resets all robosuite-internal variables
camera_names (str or list of str): name of camera to be rendered. Should either be single str if
same name is to be used for all cameras' rendering or else it should be a list of cameras to render.
:Note: At least one camera must be specified if @use_camera_obs is True.
:Note: To render all robots' cameras of a certain type (e.g.: "robotview" or "eye_in_hand"), use the
convention "all-{name}" (e.g.: "all-robotview") to automatically render all camera images from each
robot's camera list).
camera_heights (int or list of int): height of camera frame. Should either be single int if
same height is to be used for all cameras' frames or else it should be a list of the same length as
"camera names" param.
camera_widths (int or list of int): width of camera frame. Should either be single int if
same width is to be used for all cameras' frames or else it should be a list of the same length as
"camera names" param.
camera_depths (bool or list of bool): True if rendering RGB-D, and RGB otherwise. Should either be single
bool if same depth setting is to be used for all cameras or else it should be a list of the same length as
"camera names" param.
Raises:
ValueError: [Invalid number of robots specified]
ValueError: [Invalid env configuration]
ValueError: [Invalid robots for specified env configuration]
"""
def __init__(
self,
robots,
env_configuration="default",
controller_configs=None,
gripper_types="default",
initialization_noise="default",
table_full_size=(0.8, 0.8, 0.05),
table_friction=(1., 5e-3, 1e-4),
use_camera_obs=True,
use_object_obs=True,
reward_scale=1.0,
reward_shaping=False,
placement_initializer=None,
has_renderer=False,
has_offscreen_renderer=True,
render_camera="frontview",
render_collision_mesh=False,
render_visual_mesh=True,
render_gpu_device_id=-1,
control_freq=20,
horizon=1000,
ignore_done=False,
hard_reset=True,
camera_names="agentview",
camera_heights=256,
camera_widths=256,
camera_depths=False,
path=[None]
):
# settings for table top
self.table_full_size = table_full_size
self.table_friction = table_friction
self.table_offset = np.array((0, 0, 0.8))
# reward configuration
self.reward_scale = reward_scale
self.reward_shaping = reward_shaping
# whether to use ground-truth object states
self.use_object_obs = use_object_obs
# object placement initializer
self.placement_initializer = placement_initializer
super().__init__(
robots=robots,
env_configuration=env_configuration,
controller_configs=controller_configs,
mount_types="default",
gripper_types=gripper_types,
initialization_noise=initialization_noise,
use_camera_obs=use_camera_obs,
has_renderer=has_renderer,
has_offscreen_renderer=has_offscreen_renderer,
render_camera=render_camera,
render_collision_mesh=render_collision_mesh,
render_visual_mesh=render_visual_mesh,
render_gpu_device_id=render_gpu_device_id,
control_freq=control_freq,
horizon=horizon,
ignore_done=ignore_done,
hard_reset=hard_reset,
camera_names=camera_names,
camera_heights=camera_heights,
camera_widths=camera_widths,
camera_depths=camera_depths,
path=path
)
def reward(self, action=None):
"""
Reward function for the task.
Sparse un-normalized reward:
- a discrete reward of 3.0 is provided if the pot is lifted and is parallel within 30 deg to the table
Un-normalized summed components if using reward shaping:
- Reaching: in [0, 0.5], per-arm component that is proportional to the distance between each arm and its
respective pot handle, and exactly 0.5 when grasping the handle
- Note that the agent only gets the lifting reward when flipping no more than 30 degrees.
- Grasping: in {0, 0.25}, binary per-arm component awarded if the gripper is grasping its correct handle
- Lifting: in [0, 1.5], proportional to the pot's height above the table, and capped at a certain threshold
Note that the final reward is normalized and scaled by reward_scale / 3.0 as
well so that the max score is equal to reward_scale
Args:
action (np array): [NOT USED]
Returns:
float: reward value
"""
reward = 0
# check if the pot is tilted more than 30 degrees
mat = T.quat2mat(self._pot_quat)
z_unit = [0, 0, 1]
z_rotated = np.matmul(mat, z_unit)
cos_z = np.dot(z_unit, z_rotated)
cos_30 = np.cos(np.pi / 6)
direction_coef = 1 if cos_z >= cos_30 else 0
# check for goal completion: cube is higher than the table top above a margin
if self._check_success():
reward = 3.0 * direction_coef
# use a shaping reward
elif self.reward_shaping:
# lifting reward
pot_bottom_height = self.sim.data.site_xpos[self.pot_center_id][2] - self.pot.top_offset[2]
table_height = self.sim.data.site_xpos[self.table_top_id][2]
elevation = pot_bottom_height - table_height
r_lift = min(max(elevation - 0.05, 0), 0.15)
reward += 10. * direction_coef * r_lift
_gripper0_to_handle0 = self._gripper0_to_handle0
_gripper1_to_handle1 = self._gripper1_to_handle1
# gh stands for gripper-handle
# When grippers are far away, tell them to be closer
# Get contacts
(g0, g1) = (self.robots[0].gripper["right"], self.robots[0].gripper["left"]) if \
self.env_configuration == "bimanual" else (self.robots[0].gripper, self.robots[1].gripper)
_g0h_dist = np.linalg.norm(_gripper0_to_handle0)
_g1h_dist = np.linalg.norm(_gripper1_to_handle1)
# Grasping reward
if self._check_grasp(gripper=g0, object_geoms=self.pot.handle0_geoms):
reward += 0.25
# Reaching reward
reward += 0.5 * (1 - np.tanh(10.0 * _g0h_dist))
# Grasping reward
if self._check_grasp(gripper=g1, object_geoms=self.pot.handle1_geoms):
reward += 0.25
# Reaching reward
reward += 0.5 * (1 - np.tanh(10.0 * _g1h_dist))
if self.reward_scale is not None:
reward *= self.reward_scale / 3.0
return reward
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 _setup_references(self):
"""
Sets up references to important components. A reference is typically an
index or a list of indices that point to the corresponding elements
in a flatten array, which is how MuJoCo stores physical simulation data.
"""
super()._setup_references()
# Additional object references from this env
self.pot_body_id = self.sim.model.body_name2id(self.pot.root_body)
self.handle0_site_id = self.sim.model.site_name2id(self.pot.important_sites["handle0"])
self.handle1_site_id = self.sim.model.site_name2id(self.pot.important_sites["handle1"])
self.table_top_id = self.sim.model.site_name2id("table_top")
self.pot_center_id = self.sim.model.site_name2id(self.pot.important_sites["center"])
def _setup_observables(self):
"""
Sets up observables to be used for this environment. Creates object-based observables if enabled
Returns:
OrderedDict: Dictionary mapping observable names to its corresponding Observable object
"""
observables = super()._setup_observables()
# low-level object information
if self.use_object_obs:
# Get robot prefix and define observables modality
if self.env_configuration == "bimanual":
pf0 = self.robots[0].robot_model.naming_prefix + "right_"
pf1 = self.robots[0].robot_model.naming_prefix + "left_"
else:
pf0 = self.robots[0].robot_model.naming_prefix
pf1 = self.robots[1].robot_model.naming_prefix
modality = "object"
# position and rotation of object
@sensor(modality=modality)
def pot_pos(obs_cache):
return np.array(self.sim.data.body_xpos[self.pot_body_id])
@sensor(modality=modality)
def pot_quat(obs_cache):
return T.convert_quat(self.sim.data.body_xquat[self.pot_body_id], to="xyzw")
@sensor(modality=modality)
def handle0_xpos(obs_cache):
return np.array(self._handle0_xpos)
@sensor(modality=modality)
def handle1_xpos(obs_cache):
return np.array(self._handle1_xpos)
@sensor(modality=modality)
def gripper0_to_handle0(obs_cache):
return obs_cache["handle0_xpos"] - obs_cache[f"{pf0}eef_pos"] if \
"handle0_xpos" in obs_cache and f"{pf0}eef_pos" in obs_cache else np.zeros(3)
@sensor(modality=modality)
def gripper1_to_handle1(obs_cache):
return obs_cache["handle1_xpos"] - obs_cache[f"{pf1}eef_pos"] if \
"handle1_xpos" in obs_cache and f"{pf1}eef_pos" in obs_cache else np.zeros(3)
sensors = [pot_pos, pot_quat, handle0_xpos, handle1_xpos, gripper0_to_handle0, gripper1_to_handle1]
names = [s.__name__ for s in sensors]
# Create observables
for name, s in zip(names, sensors):
observables[name] = Observable(
name=name,
sensor=s,
sampling_rate=self.control_freq,
)
return observables
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():
self.sim.data.set_joint_qpos(obj.joints[0], np.concatenate([np.array(obj_pos), np.array(obj_quat)]))
def visualize(self, vis_settings):
"""
In addition to super call, visualize gripper site proportional to the distance to each handle.
Args:
vis_settings (dict): Visualization keywords mapped to T/F, determining whether that specific
component should be visualized. Should have "grippers" keyword as well as any other relevant
options specified.
"""
# Run superclass method first
super().visualize(vis_settings=vis_settings)
# Color the gripper visualization site according to its distance to each handle
if vis_settings["grippers"]:
handles = [self.pot.important_sites[f"handle{i}"] for i in range(2)]
grippers = [self.robots[0].gripper[arm] for arm in self.robots[0].arms] if \
self.env_configuration == "bimanual" else [robot.gripper for robot in self.robots]
for gripper, handle in zip(grippers, handles):
self._visualize_gripper_to_target(gripper=gripper, target=handle, target_type="site")
def _check_success(self):
"""
Check if pot is successfully lifted
Returns:
bool: True if pot is lifted
"""
pot_bottom_height = self.sim.data.site_xpos[self.pot_center_id][2] - self.pot.top_offset[2]
table_height = self.sim.data.site_xpos[self.table_top_id][2]
# cube is higher than the table top above a margin
return pot_bottom_height > table_height + 0.10
@property
def _handle0_xpos(self):
"""
Grab the position of the left (blue) hammer handle.
Returns:
np.array: (x,y,z) position of handle
"""
return self.sim.data.site_xpos[self.handle0_site_id]
@property
def _handle1_xpos(self):
"""
Grab the position of the right (green) hammer handle.
Returns:
np.array: (x,y,z) position of handle
"""
return self.sim.data.site_xpos[self.handle1_site_id]
@property
def _pot_quat(self):
"""
Grab the orientation of the pot body.
Returns:
np.array: (x,y,z,w) quaternion of the pot body
"""
return T.convert_quat(self.sim.data.body_xquat[self.pot_body_id], to="xyzw")
@property
def _gripper0_to_handle0(self):
"""
Calculate vector from the left gripper to the left pot handle.
Returns:
np.array: (dx,dy,dz) distance vector between handle and EEF0
"""
return self._handle0_xpos - self._eef0_xpos
@property
def _gripper1_to_handle1(self):
"""
Calculate vector from the right gripper to the right pot handle.
Returns:
np.array: (dx,dy,dz) distance vector between handle and EEF0
"""
return self._handle1_xpos - self._eef1_xpos
class Door(SingleArmEnv):
"""
This class corresponds to the door opening task for a single robot arm.
Args:
robots (str or list of str): Specification for specific robot arm(s) to be instantiated within this env
(e.g: "Sawyer" would generate one arm; ["Panda", "Panda", "Sawyer"] would generate three robot arms)
Note: Must be a single single-arm robot!
env_configuration (str): Specifies how to position the robots within the environment (default is "default").
For most single arm environments, this argument has no impact on the robot setup.
controller_configs (str or list of dict): If set, contains relevant controller parameters for creating a
custom controller. Else, uses the default controller for this specific task. Should either be single
dict if same controller is to be used for all robots or else it should be a list of the same length as
"robots" param
gripper_types (str or list of str): type of gripper, used to instantiate
gripper models from gripper factory. Default is "default", which is the default grippers(s) associated
with the robot(s) the 'robots' specification. None removes the gripper, and any other (valid) model
overrides the default gripper. Should either be single str if same gripper type is to be used for all
robots or else it should be a list of the same length as "robots" param
initialization_noise (dict or list of dict): Dict containing the initialization noise parameters.
The expected keys and corresponding value types are specified below:
:`'magnitude'`: The scale factor of uni-variate random noise applied to each of a robot's given initial
joint positions. Setting this value to `None` or 0.0 results in no noise being applied.
If "gaussian" type of noise is applied then this magnitude scales the standard deviation applied,
If "uniform" type of noise is applied then this magnitude sets the bounds of the sampling range
:`'type'`: Type of noise to apply. Can either specify "gaussian" or "uniform"
Should either be single dict if same noise value is to be used for all robots or else it should be a
list of the same length as "robots" param
:Note: Specifying "default" will automatically use the default noise settings.
Specifying None will automatically create the required dict with "magnitude" set to 0.0.
use_latch (bool): if True, uses a spring-loaded handle and latch to "lock" the door closed initially
Otherwise, door is instantiated with a fixed handle
use_camera_obs (bool): if True, every observation includes rendered image(s)
use_object_obs (bool): if True, include object (cube) information in
the observation.
reward_scale (None or float): Scales the normalized reward function by the amount specified.
If None, environment reward remains unnormalized
reward_shaping (bool): if True, use dense rewards.
placement_initializer (ObjectPositionSampler): if provided, will
be used to place objects on every reset, else a UniformRandomSampler
is used by default.
has_renderer (bool): If true, render the simulation state in
a viewer instead of headless mode.
has_offscreen_renderer (bool): True if using off-screen rendering
render_camera (str): Name of camera to render if `has_renderer` is True. Setting this value to 'None'
will result in the default angle being applied, which is useful as it can be dragged / panned by
the user using the mouse
render_collision_mesh (bool): True if rendering collision meshes in camera. False otherwise.
render_visual_mesh (bool): True if rendering visual meshes in camera. False otherwise.
render_gpu_device_id (int): corresponds to the GPU device id to use for offscreen rendering.
Defaults to -1, in which case the device will be inferred from environment variables
(GPUS or CUDA_VISIBLE_DEVICES).
control_freq (float): how many control signals to receive in every second. This sets the amount of
simulation time that passes between every action input.
horizon (int): Every episode lasts for exactly @horizon timesteps.
ignore_done (bool): True if never terminating the environment (ignore @horizon).
hard_reset (bool): If True, re-loads model, sim, and render object upon a reset call, else,
only calls sim.reset and resets all robosuite-internal variables
camera_names (str or list of str): name of camera to be rendered. Should either be single str if
same name is to be used for all cameras' rendering or else it should be a list of cameras to render.
:Note: At least one camera must be specified if @use_camera_obs is True.
:Note: To render all robots' cameras of a certain type (e.g.: "robotview" or "eye_in_hand"), use the
convention "all-{name}" (e.g.: "all-robotview") to automatically render all camera images from each
robot's camera list).
camera_heights (int or list of int): height of camera frame. Should either be single int if
same height is to be used for all cameras' frames or else it should be a list of the same length as
"camera names" param.
camera_widths (int or list of int): width of camera frame. Should either be single int if
same width is to be used for all cameras' frames or else it should be a list of the same length as
"camera names" param.
camera_depths (bool or list of bool): True if rendering RGB-D, and RGB otherwise. Should either be single
bool if same depth setting is to be used for all cameras or else it should be a list of the same length as
"camera names" param.
camera_segmentations (None or str or list of str or list of list of str): Camera segmentation(s) to use
for each camera. Valid options are:
`None`: no segmentation sensor used
`'instance'`: segmentation at the class-instance level
`'class'`: segmentation at the class level
`'element'`: segmentation at the per-geom level
If not None, multiple types of segmentations can be specified. A [list of str / str or None] specifies
[multiple / a single] segmentation(s) to use for all cameras. A list of list of str specifies per-camera
segmentation setting(s) to use.
Raises:
AssertionError: [Invalid number of robots specified]
"""
def __init__(
self,
robots,
env_configuration="default",
controller_configs=None,
gripper_types="default",
initialization_noise="default",
use_latch=True,
use_camera_obs=True,
use_object_obs=True,
reward_scale=1.0,
reward_shaping=False,
placement_initializer=None,
has_renderer=False,
has_offscreen_renderer=True,
render_camera="frontview",
render_collision_mesh=False,
render_visual_mesh=True,
render_gpu_device_id=-1,
control_freq=20,
horizon=1000,
ignore_done=False,
hard_reset=True,
camera_names="agentview",
camera_heights=256,
camera_widths=256,
camera_depths=False,
camera_segmentations=None, # {None, instance, class, element}
renderer="mujoco",
renderer_config=None,
):
# settings for table top (hardcoded since it's not an essential part of the environment)
self.table_full_size = (0.8, 0.3, 0.05)
self.table_offset = (-0.2, -0.35, 0.8)
# reward configuration
self.use_latch = use_latch
self.reward_scale = reward_scale
self.reward_shaping = reward_shaping
# whether to use ground-truth object states
self.use_object_obs = use_object_obs
# object placement initializer
self.placement_initializer = placement_initializer
super().__init__(
robots=robots,
env_configuration=env_configuration,
controller_configs=controller_configs,
mount_types="default",
gripper_types=gripper_types,
initialization_noise=initialization_noise,
use_camera_obs=use_camera_obs,
has_renderer=has_renderer,
has_offscreen_renderer=has_offscreen_renderer,
render_camera=render_camera,
render_collision_mesh=render_collision_mesh,
render_visual_mesh=render_visual_mesh,
render_gpu_device_id=render_gpu_device_id,
control_freq=control_freq,
horizon=horizon,
ignore_done=ignore_done,
hard_reset=hard_reset,
camera_names=camera_names,
camera_heights=camera_heights,
camera_widths=camera_widths,
camera_depths=camera_depths,
camera_segmentations=camera_segmentations,
renderer=renderer,
renderer_config=renderer_config,
)
def reward(self, action=None):
"""
Reward function for the task.
Sparse un-normalized reward:
- a discrete reward of 1.0 is provided if the door is opened
Un-normalized summed components if using reward shaping:
- Reaching: in [0, 0.25], proportional to the distance between door handle and robot arm
- Rotating: in [0, 0.25], proportional to angle rotated by door handled
- Note that this component is only relevant if the environment is using the locked door version
Note that a successfully completed task (door opened) will return 1.0 irregardless of whether the environment
is using sparse or shaped rewards
Note that the final reward is normalized and scaled by reward_scale / 1.0 as
well so that the max score is equal to reward_scale
Args:
action (np.array): [NOT USED]
Returns:
float: reward value
"""
reward = 0.0
# sparse completion reward
if self._check_success():
reward = 1.0
# else, we consider only the case if we're using shaped rewards
elif self.reward_shaping:
# Add reaching component
dist = np.linalg.norm(self._gripper_to_handle)
reaching_reward = 0.25 * (1 - np.tanh(10.0 * dist))
reward += reaching_reward
# Add rotating component if we're using a locked door
if self.use_latch:
handle_qpos = self.sim.data.qpos[self.handle_qpos_addr]
reward += np.clip(0.25 * np.abs(handle_qpos / (0.5 * np.pi)),
-0.25, 0.25)
# Scale reward if requested
if self.reward_scale is not None:
reward *= self.reward_scale / 1.0
return reward
def _load_model(self):
"""
Loads an xml model, puts it in self.model
"""
super()._load_model()
# Adjust base pose accordingly
xpos = self.robots[0].robot_model.base_xpos_offset["table"](
self.table_full_size[0])
self.robots[0].robot_model.set_base_xpos(xpos)
# load model for table top workspace
mujoco_arena = TableArena(
table_full_size=self.table_full_size,
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.5986131746834771, -4.392035683362857e-09, 1.5903500240372423
],
quat=[
0.6380177736282349, 0.3048497438430786, 0.30484986305236816,
0.6380177736282349
],
)
# initialize objects of interest
self.door = DoorObject(
name="Door",
friction=0.0,
damping=0.1,
lock=self.use_latch,
)
# Create placement initializer
if self.placement_initializer is not None:
self.placement_initializer.reset()
self.placement_initializer.add_objects(self.door)
else:
self.placement_initializer = UniformRandomSampler(
name="ObjectSampler",
mujoco_objects=self.door,
x_range=[0.07, 0.09],
y_range=[-0.01, 0.01],
rotation=(-np.pi / 2.0 - 0.25, -np.pi / 2.0),
rotation_axis="z",
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.door,
)
def _setup_references(self):
"""
Sets up references to important components. A reference is typically an
index or a list of indices that point to the corresponding elements
in a flatten array, which is how MuJoCo stores physical simulation data.
"""
super()._setup_references()
# Additional object references from this env
self.object_body_ids = dict()
self.object_body_ids["door"] = self.sim.model.body_name2id(
self.door.door_body)
self.object_body_ids["frame"] = self.sim.model.body_name2id(
self.door.frame_body)
self.object_body_ids["latch"] = self.sim.model.body_name2id(
self.door.latch_body)
self.door_handle_site_id = self.sim.model.site_name2id(
self.door.important_sites["handle"])
self.hinge_qpos_addr = self.sim.model.get_joint_qpos_addr(
self.door.joints[0])
if self.use_latch:
self.handle_qpos_addr = self.sim.model.get_joint_qpos_addr(
self.door.joints[1])
def _setup_observables(self):
"""
Sets up observables to be used for this environment. Creates object-based observables if enabled
Returns:
OrderedDict: Dictionary mapping observable names to its corresponding Observable object
"""
observables = super()._setup_observables()
# low-level object information
if self.use_object_obs:
# Get robot prefix and define observables modality
pf = self.robots[0].robot_model.naming_prefix
modality = "object"
# Define sensor callbacks
@sensor(modality=modality)
def door_pos(obs_cache):
return np.array(
self.sim.data.body_xpos[self.object_body_ids["door"]])
@sensor(modality=modality)
def handle_pos(obs_cache):
return self._handle_xpos
@sensor(modality=modality)
def door_to_eef_pos(obs_cache):
return (obs_cache["door_pos"] -
obs_cache[f"{pf}eef_pos"] if "door_pos" in obs_cache
and f"{pf}eef_pos" in obs_cache else np.zeros(3))
@sensor(modality=modality)
def handle_to_eef_pos(obs_cache):
return (obs_cache["handle_pos"] -
obs_cache[f"{pf}eef_pos"] if "handle_pos" in obs_cache
and f"{pf}eef_pos" in obs_cache else np.zeros(3))
@sensor(modality=modality)
def hinge_qpos(obs_cache):
return np.array([self.sim.data.qpos[self.hinge_qpos_addr]])
sensors = [
door_pos, handle_pos, door_to_eef_pos, handle_to_eef_pos,
hinge_qpos
]
names = [s.__name__ for s in sensors]
# Also append handle qpos if we're using a locked door version with rotatable handle
if self.use_latch:
@sensor(modality=modality)
def handle_qpos(obs_cache):
return np.array(
[self.sim.data.qpos[self.handle_qpos_addr]])
sensors.append(handle_qpos)
names.append("handle_qpos")
# Create observables
for name, s in zip(names, sensors):
observables[name] = Observable(
name=name,
sensor=s,
sampling_rate=self.control_freq,
)
return observables
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()
# We know we're only setting a single object (the door), so specifically set its pose
door_pos, door_quat, _ = object_placements[self.door.name]
door_body_id = self.sim.model.body_name2id(self.door.root_body)
self.sim.model.body_pos[door_body_id] = door_pos
self.sim.model.body_quat[door_body_id] = door_quat
def _check_success(self):
"""
Check if door has been opened.
Returns:
bool: True if door has been opened
"""
hinge_qpos = self.sim.data.qpos[self.hinge_qpos_addr]
return hinge_qpos > 0.3
def visualize(self, vis_settings):
"""
In addition to super call, visualize gripper site proportional to the distance to the door handle.
Args:
vis_settings (dict): Visualization keywords mapped to T/F, determining whether that specific
component should be visualized. Should have "grippers" keyword as well as any other relevant
options specified.
"""
# Run superclass method first
super().visualize(vis_settings=vis_settings)
# Color the gripper visualization site according to its distance to the door handle
if vis_settings["grippers"]:
self._visualize_gripper_to_target(
gripper=self.robots[0].gripper,
target=self.door.important_sites["handle"],
target_type="site")
@property
def _handle_xpos(self):
"""
Grabs the position of the door handle handle.
Returns:
np.array: Door handle (x,y,z)
"""
return self.sim.data.site_xpos[self.door_handle_site_id]
@property
def _gripper_to_handle(self):
"""
Calculates distance from the gripper to the door handle.
Returns:
np.array: (x,y,z) distance between handle and eef
"""
return self._handle_xpos - self._eef_xpos
class TwoArmHandover(TwoArmEnv):
"""
This class corresponds to the handover task for two robot arms.
Args:
robots (str or list of str): Specification for specific robot arm(s) to be instantiated within this env
(e.g: "Sawyer" would generate one arm; ["Panda", "Panda", "Sawyer"] would generate three robot arms)
Note: Must be either 2 single single-arm robots or 1 bimanual robot!
env_configuration (str): Specifies how to position the robots within the environment. Can be either:
:`'bimanual'`: Only applicable for bimanual robot setups. Sets up the (single) bimanual robot on the -x
side of the table
:`'single-arm-parallel'`: Only applicable for multi single arm setups. Sets up the (two) single armed
robots next to each other on the -x side of the table
:`'single-arm-opposed'`: Only applicable for multi single arm setups. Sets up the (two) single armed
robots opposed from each others on the opposite +/-y sides of the table.
Note that "default" corresponds to either "bimanual" if a bimanual robot is used or "single-arm-opposed" if two
single-arm robots are used.
controller_configs (str or list of dict): If set, contains relevant controller parameters for creating a
custom controller. Else, uses the default controller for this specific task. Should either be single
dict if same controller is to be used for all robots or else it should be a list of the same length as
"robots" param
gripper_types (str or list of str): type of gripper, used to instantiate
gripper models from gripper factory. Default is "default", which is the default grippers(s) associated
with the robot(s) the 'robots' specification. None removes the gripper, and any other (valid) model
overrides the default gripper. Should either be single str if same gripper type is to be used for all
robots or else it should be a list of the same length as "robots" param
initialization_noise (dict or list of dict): Dict containing the initialization noise parameters.
The expected keys and corresponding value types are specified below:
:`'magnitude'`: The scale factor of uni-variate random noise applied to each of a robot's given initial
joint positions. Setting this value to `None` or 0.0 results in no noise being applied.
If "gaussian" type of noise is applied then this magnitude scales the standard deviation applied,
If "uniform" type of noise is applied then this magnitude sets the bounds of the sampling range
:`'type'`: Type of noise to apply. Can either specify "gaussian" or "uniform"
Should either be single dict if same noise value is to be used for all robots or else it should be a
list of the same length as "robots" param
:Note: Specifying "default" will automatically use the default noise settings.
Specifying None will automatically create the required dict with "magnitude" set to 0.0.
prehensile (bool): If true, handover object starts on the table. Else, the object starts in Arm0's gripper
table_full_size (3-tuple): x, y, and z dimensions of the table.
table_friction (3-tuple): the three mujoco friction parameters for
the table.
use_camera_obs (bool): if True, every observation includes rendered image(s)
use_object_obs (bool): if True, include object (cube) information in
the observation.
reward_scale (None or float): Scales the normalized reward function by the amount specified.
If None, environment reward remains unnormalized
reward_shaping (bool): if True, use dense rewards.
placement_initializer (ObjectPositionSampler): if provided, will
be used to place objects on every reset, else a UniformRandomSampler
is used by default.
has_renderer (bool): If true, render the simulation state in
a viewer instead of headless mode.
has_offscreen_renderer (bool): True if using off-screen rendering
render_camera (str): Name of camera to render if `has_renderer` is True. Setting this value to 'None'
will result in the default angle being applied, which is useful as it can be dragged / panned by
the user using the mouse
render_collision_mesh (bool): True if rendering collision meshes in camera. False otherwise.
render_visual_mesh (bool): True if rendering visual meshes in camera. False otherwise.
render_gpu_device_id (int): corresponds to the GPU device id to use for offscreen rendering.
Defaults to -1, in which case the device will be inferred from environment variables
(GPUS or CUDA_VISIBLE_DEVICES).
control_freq (float): how many control signals to receive in every second. This sets the amount of
simulation time that passes between every action input.
horizon (int): Every episode lasts for exactly @horizon timesteps.
ignore_done (bool): True if never terminating the environment (ignore @horizon).
hard_reset (bool): If True, re-loads model, sim, and render object upon a reset call, else,
only calls sim.reset and resets all robosuite-internal variables
camera_names (str or list of str): name of camera to be rendered. Should either be single str if
same name is to be used for all cameras' rendering or else it should be a list of cameras to render.
:Note: At least one camera must be specified if @use_camera_obs is True.
:Note: To render all robots' cameras of a certain type (e.g.: "robotview" or "eye_in_hand"), use the
convention "all-{name}" (e.g.: "all-robotview") to automatically render all camera images from each
robot's camera list).
camera_heights (int or list of int): height of camera frame. Should either be single int if
same height is to be used for all cameras' frames or else it should be a list of the same length as
"camera names" param.
camera_widths (int or list of int): width of camera frame. Should either be single int if
same width is to be used for all cameras' frames or else it should be a list of the same length as
"camera names" param.
camera_depths (bool or list of bool): True if rendering RGB-D, and RGB otherwise. Should either be single
bool if same depth setting is to be used for all cameras or else it should be a list of the same length as
"camera names" param.
Raises:
ValueError: [Invalid number of robots specified]
ValueError: [Invalid env configuration]
ValueError: [Invalid robots for specified env configuration]
"""
def __init__(self,
robots,
env_configuration="default",
controller_configs=None,
gripper_types="default",
initialization_noise="default",
prehensile=True,
table_full_size=(0.8, 1.2, 0.05),
table_friction=(1., 5e-3, 1e-4),
use_camera_obs=True,
use_object_obs=True,
reward_scale=1.0,
reward_shaping=False,
placement_initializer=None,
has_renderer=False,
has_offscreen_renderer=True,
render_camera="frontview",
render_collision_mesh=False,
render_visual_mesh=True,
render_gpu_device_id=-1,
control_freq=20,
horizon=1000,
ignore_done=False,
hard_reset=True,
camera_names="agentview",
camera_heights=256,
camera_widths=256,
camera_depths=False,
path=[None]):
# Task settings
self.prehensile = prehensile
# settings for table top
self.table_full_size = table_full_size
self.table_true_size = list(table_full_size)
self.table_true_size[
1] *= 0.25 # true size will only be partially wide
self.table_friction = table_friction
self.table_offset = [0, self.table_full_size[1] * (-3 / 8), 0.8]
# reward configuration
self.reward_scale = reward_scale
self.reward_shaping = reward_shaping
self.height_threshold = 0.1 # threshold above the table surface which the hammer is considered lifted
# whether to use ground-truth object states
self.use_object_obs = use_object_obs
# object placement initializer
self.placement_initializer = placement_initializer
super().__init__(robots=robots,
env_configuration=env_configuration,
controller_configs=controller_configs,
mount_types="default",
gripper_types=gripper_types,
initialization_noise=initialization_noise,
use_camera_obs=use_camera_obs,
has_renderer=has_renderer,
has_offscreen_renderer=has_offscreen_renderer,
render_camera=render_camera,
render_collision_mesh=render_collision_mesh,
render_visual_mesh=render_visual_mesh,
render_gpu_device_id=render_gpu_device_id,
control_freq=control_freq,
horizon=horizon,
ignore_done=ignore_done,
hard_reset=hard_reset,
camera_names=camera_names,
camera_heights=camera_heights,
camera_widths=camera_widths,
camera_depths=camera_depths,
path=path)
def reward(self, action=None):
"""
Reward function for the task.
Sparse un-normalized reward:
- a discrete reward of 2.0 is provided when only Arm 1 is gripping the handle and has the handle
lifted above a certain threshold
Un-normalized max-wise components if using reward shaping:
- Arm0 Reaching: (1) in [0, 0.25] proportional to the distance between Arm 0 and the handle
- Arm0 Grasping: (2) in {0, 0.5}, nonzero if Arm 0 is gripping the hammer (any part).
- Arm0 Lifting: (3) in {0, 1.0}, nonzero if Arm 0 lifts the handle from the table past a certain threshold
- Arm0 Hovering: (4) in {0, [1.0, 1.25]}, nonzero only if Arm0 is actively lifting the hammer, and is
proportional to the distance between the handle and Arm 1
conditioned on the handle being lifted from the table and being grasped by Arm 0
- Mutual Grasping: (5) in {0, 1.5}, nonzero if both Arm 0 and Arm 1 are gripping the hammer (Arm 1 must be
gripping the handle) while lifted above the table
- Handover: (6) in {0, 2.0}, nonzero when only Arm 1 is gripping the handle and has the handle
lifted above the table
Note that the final reward is normalized and scaled by reward_scale / 2.0 as
well so that the max score is equal to reward_scale
Args:
action (np array): [NOT USED]
Returns:
float: reward value
"""
# Initialize reward
reward = 0
# use a shaping reward if specified
if self.reward_shaping:
# Grab relevant parameters
arm0_grasp_any, arm1_grasp_handle, hammer_height, table_height = self._get_task_info(
)
# First, we'll consider the cases if the hammer is lifted above the threshold (step 3 - 6)
if hammer_height - table_height > self.height_threshold:
# Split cases depending on whether arm1 is currently grasping the handle or not
if arm1_grasp_handle:
# Check if arm0 is grasping
if arm0_grasp_any:
# This is step 5
reward = 1.5
else:
# This is step 6 (completed task!)
reward = 2.0
# This is the case where only arm0 is grasping (step 2-3)
else:
reward = 1.0
# Add in up to 0.25 based on distance between handle and arm1
dist = np.linalg.norm(self._gripper_1_to_handle)
reaching_reward = 0.25 * (1 - np.tanh(1.0 * dist))
reward += reaching_reward
# Else, the hammer is still on the ground ):
else:
# Split cases depending on whether arm0 is currently grasping the handle or not
if arm0_grasp_any:
# This is step 2
reward = 0.5
else:
# This is step 1, we want to encourage arm0 to reach for the handle
dist = np.linalg.norm(self._gripper_0_to_handle)
reaching_reward = 0.25 * (1 - np.tanh(1.0 * dist))
reward = reaching_reward
# Else this is the sparse reward setting
else:
# Provide reward if only Arm 1 is grasping the hammer and the handle lifted above the pre-defined threshold
if self._check_success():
reward = 2.0
if self.reward_scale is not None:
reward *= self.reward_scale / 2.0
return reward
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 _setup_references(self):
"""
Sets up references to important components. A reference is typically an
index or a list of indices that point to the corresponding elements
in a flatten array, which is how MuJoCo stores physical simulation data.
"""
super()._setup_references()
# Hammer object references from this env
self.hammer_body_id = self.sim.model.body_name2id(
self.hammer.root_body)
self.hammer_handle_geom_id = self.sim.model.geom_name2id(
self.hammer.handle_geoms[0])
# General env references
self.table_top_id = self.sim.model.site_name2id("table_top")
def _setup_observables(self):
"""
Sets up observables to be used for this environment. Creates object-based observables if enabled
Returns:
OrderedDict: Dictionary mapping observable names to its corresponding Observable object
"""
observables = super()._setup_observables()
# low-level object information
if self.use_object_obs:
# Get robot prefix and define observables modality
if self.env_configuration == "bimanual":
pf0 = self.robots[0].robot_model.naming_prefix + "right_"
pf1 = self.robots[0].robot_model.naming_prefix + "left_"
else:
pf0 = self.robots[0].robot_model.naming_prefix
pf1 = self.robots[1].robot_model.naming_prefix
modality = "object"
# position and rotation of hammer
@sensor(modality=modality)
def hammer_pos(obs_cache):
return np.array(self._hammer_pos)
@sensor(modality=modality)
def hammer_quat(obs_cache):
return np.array(self._hammer_quat)
@sensor(modality=modality)
def handle_xpos(obs_cache):
return np.array(self._handle_xpos)
@sensor(modality=modality)
def gripper0_to_handle(obs_cache):
return obs_cache["handle_xpos"] - obs_cache[f"{pf0}eef_pos"] if \
"handle_xpos" in obs_cache and f"{pf0}eef_pos" in obs_cache else np.zeros(3)
@sensor(modality=modality)
def gripper1_to_handle(obs_cache):
return obs_cache["handle_xpos"] - obs_cache[f"{pf1}eef_pos"] if \
"handle_xpos" in obs_cache and f"{pf1}eef_pos" in obs_cache else np.zeros(3)
sensors = [
hammer_pos, hammer_quat, handle_xpos, gripper0_to_handle,
gripper1_to_handle
]
names = [s.__name__ for s in sensors]
# Create observables
for name, s in zip(names, sensors):
observables[name] = Observable(
name=name,
sensor=s,
sampling_rate=self.control_freq,
)
return observables
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_task_info(self):
"""
Helper function that grabs the current relevant locations of objects of interest within the environment
Returns:
4-tuple:
- (bool) True if Arm0 is grasping any part of the hammer
- (bool) True if Arm1 is grasping the hammer handle
- (float) Height of the hammer body
- (float) Height of the table surface
"""
# Get height of hammer and table and define height threshold
hammer_angle_offset = (self.hammer.handle_length / 2 +
2 * self.hammer.head_halfsize) * np.sin(
self._hammer_angle)
hammer_height = self.sim.data.geom_xpos[self.hammer_handle_geom_id][2]\
- self.hammer.top_offset[2]\
- hammer_angle_offset
table_height = self.sim.data.site_xpos[self.table_top_id][2]
# Check if any Arm's gripper is grasping the hammer handle
(g0, g1) = (self.robots[0].gripper["right"], self.robots[0].gripper["left"]) if \
self.env_configuration == "bimanual" else (self.robots[0].gripper, self.robots[1].gripper)
arm0_grasp_any = self._check_grasp(gripper=g0,
object_geoms=self.hammer)
arm1_grasp_handle = self._check_grasp(
gripper=g1, object_geoms=self.hammer.handle_geoms)
# Return all relevant values
return arm0_grasp_any, arm1_grasp_handle, hammer_height, table_height
def _check_success(self):
"""
Check if hammer is successfully handed off
Returns:
bool: True if handover has been completed
"""
# Grab relevant params
arm0_grasp_any, arm1_grasp_handle, hammer_height, table_height = self._get_task_info(
)
return \
True if \
arm1_grasp_handle and \
not arm0_grasp_any and \
hammer_height - table_height > self.height_threshold \
else False
@property
def _handle_xpos(self):
"""
Grab the position of the hammer handle.
Returns:
np.array: (x,y,z) position of handle
"""
return self.sim.data.geom_xpos[self.hammer_handle_geom_id]
@property
def _hammer_pos(self):
"""
Grab the position of the hammer body.
Returns:
np.array: (x,y,z) position of body
"""
return np.array(self.sim.data.body_xpos[self.hammer_body_id])
@property
def _hammer_quat(self):
"""
Grab the orientation of the hammer body.
Returns:
np.array: (x,y,z,w) quaternion of the hammer body
"""
return T.convert_quat(self.sim.data.body_xquat[self.hammer_body_id],
to="xyzw")
@property
def _hammer_angle(self):
"""
Calculate the angle of hammer with the ground, relative to it resting horizontally
Returns:
float: angle in radians
"""
mat = T.quat2mat(self._hammer_quat)
z_unit = [0, 0, 1]
z_rotated = np.matmul(mat, z_unit)
return np.pi / 2 - np.arccos(np.dot(z_unit, z_rotated))
@property
def _gripper_0_to_handle(self):
"""
Calculate vector from the left gripper to the hammer handle.
Returns:
np.array: (dx,dy,dz) distance vector between handle and EEF0
"""
return self._handle_xpos - self._eef0_xpos
@property
def _gripper_1_to_handle(self):
"""
Calculate vector from the right gripper to the hammer handle.
Returns:
np.array: (dx,dy,dz) distance vector between handle and EEF1
"""
return self._handle_xpos - self._eef1_xpos
class FetchPush(SingleArmEnv):
"""
This class corresponds to the fetchpush task for a single robot arm.
Args:
robots (str or list of str): Specification for specific robot arm(s) to be instantiated within this env
(e.g: "Sawyer" would generate one arm; ["Panda", "Panda", "Sawyer"] would generate three robot arms)
Note: Must be a single single-arm robot!
env_configuration (str): Specifies how to position the robots within the environment (default is "default").
For most single arm environments, this argument has no impact on the robot setup.
controller_configs (str or list of dict): If set, contains relevant controller parameters for creating a
custom controller. Else, uses the default controller for this specific task. Should either be single
dict if same controller is to be used for all robots or else it should be a list of the same length as
"robots" param
gripper_types (str or list of str): type of gripper, used to instantiate
gripper models from gripper factory. Default is "default", which is the default grippers(s) associated
with the robot(s) the 'robots' specification. None removes the gripper, and any other (valid) model
overrides the default gripper. Should either be single str if same gripper type is to be used for all
robots or else it should be a list of the same length as "robots" param
initialization_noise (dict or list of dict): Dict containing the initialization noise parameters.
The expected keys and corresponding value types are specified below:
:`'magnitude'`: The scale factor of uni-variate random noise applied to each of a robot's given initial
joint positions. Setting this value to `None` or 0.0 results in no noise being applied.
If "gaussian" type of noise is applied then this magnitude scales the standard deviation applied,
If "uniform" type of noise is applied then this magnitude sets the bounds of the sampling range
:`'type'`: Type of noise to apply. Can either specify "gaussian" or "uniform"
Should either be single dict if same noise value is to be used for all robots or else it should be a
list of the same length as "robots" param
:Note: Specifying "default" will automatically use the default noise settings.
Specifying None will automatically create the required dict with "magnitude" set to 0.0.
table_full_size (3-tuple): x, y, and z dimensions of the table.
table_friction (3-tuple): the three mujoco friction parameters for
the table.
use_camera_obs (bool): if True, every observation includes rendered image(s)
use_object_obs (bool): if True, include object (cube) information in
the observation.
reward_scale (None or float): Scales the normalized reward function by the amount specified.
If None, environment reward remains unnormalized
reward_shaping (bool): if True, use dense rewards.
placement_initializer (ObjectPositionSampler): if provided, will
be used to place objects on every reset, else a UniformRandomSampler
is used by default.
has_renderer (bool): If true, render the simulation state in
a viewer instead of headless mode.
has_offscreen_renderer (bool): True if using off-screen rendering
render_camera (str): Name of camera to render if `has_renderer` is True. Setting this value to 'None'
will result in the default angle being applied, which is useful as it can be dragged / panned by
the user using the mouse
render_collision_mesh (bool): True if rendering collision meshes in camera. False otherwise.
render_visual_mesh (bool): True if rendering visual meshes in camera. False otherwise.
render_gpu_device_id (int): corresponds to the GPU device id to use for offscreen rendering.
Defaults to -1, in which case the device will be inferred from environment variables
(GPUS or CUDA_VISIBLE_DEVICES).
control_freq (float): how many control signals to receive in every second. This sets the amount of
simulation time that passes between every action input.
horizon (int): Every episode lasts for exactly @horizon timesteps.
ignore_done (bool): True if never terminating the environment (ignore @horizon).
hard_reset (bool): If True, re-loads model, sim, and render object upon a reset call, else,
only calls sim.reset and resets all robosuite-internal variables
camera_names (str or list of str): name of camera to be rendered. Should either be single str if
same name is to be used for all cameras' rendering or else it should be a list of cameras to render.
:Note: At least one camera must be specified if @use_camera_obs is True.
:Note: To render all robots' cameras of a certain type (e.g.: "robotview" or "eye_in_hand"), use the
convention "all-{name}" (e.g.: "all-robotview") to automatically render all camera images from each
robot's camera list).
camera_heights (int or list of int): height of camera frame. Should either be single int if
same height is to be used for all cameras' frames or else it should be a list of the same length as
"camera names" param.
camera_widths (int or list of int): width of camera frame. Should either be single int if
same width is to be used for all cameras' frames or else it should be a list of the same length as
"camera names" param.
camera_depths (bool or list of bool): True if rendering RGB-D, and RGB otherwise. Should either be single
bool if same depth setting is to be used for all cameras or else it should be a list of the same length as
"camera names" param.
Raises:
AssertionError: [Invalid number of robots specified]
"""
def __init__(
self,
robots,
env_configuration="default",
controller_configs=None,
gripper_types="default",
initialization_noise="default",
table_full_size=(0.8, 0.8, 0.05),
table_friction=(1., 5e-3, 1e-4),
use_camera_obs=True,
use_object_obs=True,
reward_scale=1.0,
reward_shaping=False,
placement_initializer=None,
has_renderer=False,
has_offscreen_renderer=True,
render_camera="frontview",
render_collision_mesh=False,
render_visual_mesh=True,
render_gpu_device_id=-1,
control_freq=20,
horizon=1000,
ignore_done=False,
hard_reset=True,
camera_names="agentview",
camera_heights=256,
camera_widths=256,
camera_depths=False,
):
# settings for table top
self.table_full_size = table_full_size
self.table_friction = table_friction
self.table_offset = np.array((0, 0, 0.8))
# reward configuration
self.reward_scale = reward_scale
self.reward_shaping = reward_shaping
# whether to use ground-truth object states
self.use_object_obs = use_object_obs
# object placement initializer
self.placement_initializer = placement_initializer
super().__init__(
robots=robots,
env_configuration=env_configuration,
controller_configs=controller_configs,
mount_types="default",
gripper_types=gripper_types,
initialization_noise=initialization_noise,
use_camera_obs=use_camera_obs,
has_renderer=has_renderer,
has_offscreen_renderer=has_offscreen_renderer,
render_camera=render_camera,
render_collision_mesh=render_collision_mesh,
render_visual_mesh=render_visual_mesh,
render_gpu_device_id=render_gpu_device_id,
control_freq=control_freq,
horizon=horizon,
ignore_done=ignore_done,
hard_reset=hard_reset,
camera_names=camera_names,
camera_heights=camera_heights,
camera_widths=camera_widths,
camera_depths=camera_depths,
)
def reward(self, action=None):
"""
Reward function for the task.
Args:
action (np array): [NOT USED]
Returns:
float: reward value
"""
reward = .0
cube_pos = np.array(self.sim.data.body_xpos[self.cube_body_id])
goal_pos = np.array(self.sim.data.body_xpos[self.goal_body_id])
gripper_site_pos = self.sim.data.site_xpos[self.robots[0].eef_site_id]
cube_to_goal_dist = np.linalg.norm(cube_pos - goal_pos)
gripper_to_cube_dist = np.linalg.norm(gripper_site_pos - cube_pos)
if self.reward_shaping:
# Return large reward if cube has been moved to goal
if self._check_success():
reward += 2.25
# Reaching reward
reward += 1 - np.tanh(10.0 * gripper_to_cube_dist)
if self.check_contact(self.robots[0].gripper, self.cube):
reward += 0.5 # Reward for touching cube
reward += 1.5 - 1.5 * np.tanh(
10.0 * cube_to_goal_dist) # Reward for pushing cube
else:
reward -= float(not self._check_success())
return reward
def _load_model(self):
"""
Loads an xml model, puts it in self.model
"""
super()._load_model()
# Adjust base pose accordingly
xpos = self.robots[0].robot_model.base_xpos_offset["table"](
self.table_full_size[0])
self.robots[0].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.cube = BoxObject(
name="cube",
size_min=[0.020, 0.020, 0.020], # [0.015, 0.015, 0.015],
size_max=[0.022, 0.022, 0.022], # [0.018, 0.018, 0.018])
rgba=[1, 0, 0, 1],
)
self.goal = BoxObject(
name="goal",
size_min=[0.005, 0.005, 0.005],
size_max=[0.005, 0.005, 0.005],
rgba=[0, 0, 1, 1],
density=11342, # Density of lead such that the goal is hard to move
)
# Create placement initializer
if self.placement_initializer is not None:
self.placement_initializer.reset()
self.placement_initializer.add_objects([self.cube, self.goal])
else:
self.placement_initializer = UniformRandomSampler(
name="ObjectSampler",
mujoco_objects=[self.cube, self.goal],
x_range=[-0.20, 0.20],
y_range=[-0.20, 0.20],
rotation=None,
ensure_object_boundary_in_range=False,
ensure_valid_placement=True,
reference_pos=self.table_offset,
z_offset=0.01,
)
# 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.cube, self.goal],
)
def _setup_references(self):
"""
Sets up references to important components. A reference is typically an
index or a list of indices that point to the corresponding elements
in a flatten array, which is how MuJoCo stores physical simulation data.
"""
super()._setup_references()
# Additional object references from this env
self.cube_body_id = self.sim.model.body_name2id(self.cube.root_body)
self.goal_body_id = self.sim.model.body_name2id(self.goal.root_body)
def _setup_observables(self):
"""
Sets up observables to be used for this environment. Creates object-based observables if enabled
Returns:
OrderedDict: Dictionary mapping observable names to its corresponding Observable object
"""
observables = super()._setup_observables()
# low-level object information
if self.use_object_obs:
# Get robot prefix and define observables modality
pf = self.robots[0].robot_model.naming_prefix
modality = "object"
# cube-related observables
@sensor(modality=modality)
def cube_pos(obs_cache):
return np.array(self.sim.data.body_xpos[self.cube_body_id])
@sensor(modality=modality)
def cube_quat(obs_cache):
return convert_quat(np.array(
self.sim.data.body_xquat[self.cube_body_id]),
to="xyzw")
@sensor(modality=modality)
def gripper_to_cube_pos(obs_cache):
return obs_cache[f"{pf}eef_pos"] - obs_cache["cube_pos"] if \
f"{pf}eef_pos" in obs_cache and "cube_pos" in obs_cache else np.zeros(3)
# goal-related observables
@sensor(modality=modality)
def goal_pos(obs_cache):
return np.array(self.sim.data.body_xpos[self.goal_body_id])
@sensor(modality=modality)
def cube_to_goal_pos(obs_cache):
return obs_cache["cube_pos"] - obs_cache["goal_pos"] if \
"cube_pos" in obs_cache and "goal_pos" in obs_cache else np.zeros(3)
sensors = [
cube_pos, cube_quat, gripper_to_cube_pos, goal_pos,
cube_to_goal_pos
]
names = [s.__name__ for s in sensors]
# Create observables
for name, s in zip(names, sensors):
observables[name] = Observable(
name=name,
sensor=s,
sampling_rate=self.control_freq,
)
return observables
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():
self.sim.data.set_joint_qpos(
obj.joints[0],
np.concatenate([np.array(obj_pos),
np.array(obj_quat)]))
def visualize(self, vis_settings):
"""
In addition to super call, visualize gripper site proportional to the distance to the cube.
Args:
vis_settings (dict): Visualization keywords mapped to T/F, determining whether that specific
component should be visualized. Should have "grippers" keyword as well as any other relevant
options specified.
"""
# Run superclass method first
super().visualize(vis_settings=vis_settings)
# Color the gripper visualization site according to its distance to the cube
if vis_settings["grippers"]:
self._visualize_gripper_to_target(gripper=self.robots[0].gripper,
target=self.cube)
def _check_success(self):
"""
Check if cube has been pushed to goal.
Returns:
bool: True if cube has been pushed to goal
"""
cube_pos = np.array(self.sim.data.body_xpos[self.cube_body_id])
goal_pos = np.array(self.sim.data.body_xpos[self.goal_body_id])
cube_to_goal_dist = np.linalg.norm(goal_pos - cube_pos)
# cube is within a margin away from goal
return cube_to_goal_dist < 0.05
class Stack(SingleArmEnv):
"""
This class corresponds to the stacking task for a single robot arm.
Args:
robots (str or list of str): Specification for specific robot arm(s) to be instantiated within this env
(e.g: "Sawyer" would generate one arm; ["Panda", "Panda", "Sawyer"] would generate three robot arms)
Note: Must be a single single-arm robot!
env_configuration (str): Specifies how to position the robots within the environment (default is "default").
For most single arm environments, this argument has no impact on the robot setup.
controller_configs (str or list of dict): If set, contains relevant controller parameters for creating a
custom controller. Else, uses the default controller for this specific task. Should either be single
dict if same controller is to be used for all robots or else it should be a list of the same length as
"robots" param
gripper_types (str or list of str): type of gripper, used to instantiate
gripper models from gripper factory. Default is "default", which is the default grippers(s) associated
with the robot(s) the 'robots' specification. None removes the gripper, and any other (valid) model
overrides the default gripper. Should either be single str if same gripper type is to be used for all
robots or else it should be a list of the same length as "robots" param
initialization_noise (dict or list of dict): Dict containing the initialization noise parameters.
The expected keys and corresponding value types are specified below:
:`'magnitude'`: The scale factor of uni-variate random noise applied to each of a robot's given initial
joint positions. Setting this value to `None` or 0.0 results in no noise being applied.
If "gaussian" type of noise is applied then this magnitude scales the standard deviation applied,
If "uniform" type of noise is applied then this magnitude sets the bounds of the sampling range
:`'type'`: Type of noise to apply. Can either specify "gaussian" or "uniform"
Should either be single dict if same noise value is to be used for all robots or else it should be a
list of the same length as "robots" param
:Note: Specifying "default" will automatically use the default noise settings.
Specifying None will automatically create the required dict with "magnitude" set to 0.0.
table_full_size (3-tuple): x, y, and z dimensions of the table.
table_friction (3-tuple): the three mujoco friction parameters for
the table.
use_camera_obs (bool): if True, every observation includes rendered image(s)
use_object_obs (bool): if True, include object (cube) information in
the observation.
reward_scale (None or float): Scales the normalized reward function by the amount specified.
If None, environment reward remains unnormalized
reward_shaping (bool): if True, use dense rewards.
placement_initializer (ObjectPositionSampler): if provided, will
be used to place objects on every reset, else a UniformRandomSampler
is used by default.
has_renderer (bool): If true, render the simulation state in
a viewer instead of headless mode.
has_offscreen_renderer (bool): True if using off-screen rendering
render_camera (str): Name of camera to render if `has_renderer` is True. Setting this value to 'None'
will result in the default angle being applied, which is useful as it can be dragged / panned by
the user using the mouse
render_collision_mesh (bool): True if rendering collision meshes in camera. False otherwise.
render_visual_mesh (bool): True if rendering visual meshes in camera. False otherwise.
render_gpu_device_id (int): corresponds to the GPU device id to use for offscreen rendering.
Defaults to -1, in which case the device will be inferred from environment variables
(GPUS or CUDA_VISIBLE_DEVICES).
control_freq (float): how many control signals to receive in every second. This sets the amount of
simulation time that passes between every action input.
horizon (int): Every episode lasts for exactly @horizon timesteps.
ignore_done (bool): True if never terminating the environment (ignore @horizon).
hard_reset (bool): If True, re-loads model, sim, and render object upon a reset call, else,
only calls sim.reset and resets all robosuite-internal variables
camera_names (str or list of str): name of camera to be rendered. Should either be single str if
same name is to be used for all cameras' rendering or else it should be a list of cameras to render.
:Note: At least one camera must be specified if @use_camera_obs is True.
:Note: To render all robots' cameras of a certain type (e.g.: "robotview" or "eye_in_hand"), use the
convention "all-{name}" (e.g.: "all-robotview") to automatically render all camera images from each
robot's camera list).
camera_heights (int or list of int): height of camera frame. Should either be single int if
same height is to be used for all cameras' frames or else it should be a list of the same length as
"camera names" param.
camera_widths (int or list of int): width of camera frame. Should either be single int if
same width is to be used for all cameras' frames or else it should be a list of the same length as
"camera names" param.
camera_depths (bool or list of bool): True if rendering RGB-D, and RGB otherwise. Should either be single
bool if same depth setting is to be used for all cameras or else it should be a list of the same length as
"camera names" param.
Raises:
AssertionError: [Invalid number of robots specified]
"""
def __init__(self,
robots,
env_configuration="default",
controller_configs=None,
gripper_types="default",
initialization_noise="default",
table_full_size=(0.8, 0.8, 0.05),
table_friction=(1., 5e-3, 1e-4),
use_camera_obs=True,
use_object_obs=True,
reward_scale=1.0,
reward_shaping=False,
placement_initializer=None,
has_renderer=False,
has_offscreen_renderer=True,
render_camera="frontview",
render_collision_mesh=False,
render_visual_mesh=True,
render_gpu_device_id=-1,
control_freq=20,
horizon=1000,
ignore_done=False,
hard_reset=True,
camera_names="agentview",
camera_heights=256,
camera_widths=256,
camera_depths=False,
path=[None]):
# settings for table top
self.table_full_size = table_full_size
self.table_friction = table_friction
self.table_offset = np.array((0, 0, 0.8))
# reward configuration
self.reward_scale = reward_scale
self.reward_shaping = reward_shaping
# whether to use ground-truth object states
self.use_object_obs = use_object_obs
# object placement initializer
self.placement_initializer = placement_initializer
self.lift_times = 0
self.hold_times = 0
self.align_reward = 0
self.stack_times = 0
self.get_info = True
super().__init__(robots=robots,
env_configuration=env_configuration,
controller_configs=controller_configs,
mount_types="default",
gripper_types=gripper_types,
initialization_noise=initialization_noise,
use_camera_obs=use_camera_obs,
has_renderer=has_renderer,
has_offscreen_renderer=has_offscreen_renderer,
render_camera=render_camera,
render_collision_mesh=render_collision_mesh,
render_visual_mesh=render_visual_mesh,
render_gpu_device_id=render_gpu_device_id,
control_freq=control_freq,
horizon=horizon,
ignore_done=ignore_done,
hard_reset=hard_reset,
camera_names=camera_names,
camera_heights=camera_heights,
camera_widths=camera_widths,
camera_depths=camera_depths,
path=path)
def reward(self, action):
"""
Reward function for the task.
Sparse un-normalized reward:
- a discrete reward of 2.0 is provided if the red block is stacked on the green block
Un-normalized components if using reward shaping:
- Reaching: in [0, 0.25], to encourage the arm to reach the cube
- Grasping: in {0, 0.25}, non-zero if arm is grasping the cube
- Lifting: in {0, 1}, non-zero if arm has lifted the cube
- Aligning: in [0, 0.5], encourages aligning one cube over the other
- Stacking: in {0, 2}, non-zero if cube is stacked on other cube
The reward is max over the following:
- Reaching + Grasping
- Lifting + Aligning
- Stacking
The sparse reward only consists of the stacking component.
Note that the final reward is normalized and scaled by
reward_scale / 2.0 as well so that the max score is equal to reward_scale
Args:
action (np array): [NOT USED]
Returns:
float: reward value
"""
r_reach, r_lift, r_stack = self.staged_rewards()
if self.reward_shaping:
reward = max(r_reach, r_lift, r_stack)
else:
reward = 2.0 if r_stack > 0 else 0.0
if self.reward_scale is not None:
reward *= self.reward_scale / 2.0
return reward
def staged_rewards(self):
"""
Helper function to calculate staged rewards based on current physical states.
Returns:
3-tuple:
- (float): reward for reaching and grasping
- (float): reward for lifting and aligning
- (float): reward for stacking
"""
# reaching is successful when the gripper site is close to the center of the cube
cubeA_pos = self.sim.data.body_xpos[self.cubeA_body_id]
cubeB_pos = self.sim.data.body_xpos[self.cubeB_body_id]
gripper_site_pos = self.sim.data.site_xpos[self.robots[0].eef_site_id]
dist = np.linalg.norm(gripper_site_pos - cubeA_pos)
r_reach = (1 - np.tanh(10.0 * dist)) * 0.25
# grasping reward
grasping_cubeA = self._check_grasp(gripper=self.robots[0].gripper,
object_geoms=self.cubeA)
if grasping_cubeA:
r_reach += 0.25
self.hold_times += 1
# lifting is successful when the cube is above the table top by a margin
cubeA_height = cubeA_pos[2]
table_height = self.table_offset[2]
cubeA_lifted = cubeA_height > table_height + 0.04
r_lift = 1.0 if cubeA_lifted else 0.0
if cubeA_lifted:
self.lift_times += 1
# Aligning is successful when cubeA is right above cubeB
if cubeA_lifted:
horiz_dist = np.linalg.norm(
np.array(cubeA_pos[:2]) - np.array(cubeB_pos[:2]))
r_lift += 0.5 * (1 - np.tanh(horiz_dist))
self.align_reward += 0.5 * (1 - np.tanh(horiz_dist))
# stacking is successful when the block is lifted and the gripper is not holding the object
r_stack = 0
cubeA_touching_cubeB = self.check_contact(self.cubeA, self.cubeB)
if not grasping_cubeA and r_lift > 0 and cubeA_touching_cubeB:
r_stack = 2.0
self.stack_times += 1
return r_reach, r_lift, r_stack
def _load_model(self):
"""
Loads an xml model, puts it in self.model
"""
super()._load_model()
# Adjust base pose accordingly
xpos = self.robots[0].robot_model.base_xpos_offset["table"](
self.table_full_size[0])
self.robots[0].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
tex_attrib = {
"type": "cube",
}
mat_attrib = {
"texrepeat": "1 1",
"specular": "0.4",
"shininess": "0.1",
}
redwood = CustomMaterial(
texture="WoodRed",
tex_name="redwood",
mat_name="redwood_mat",
tex_attrib=tex_attrib,
mat_attrib=mat_attrib,
)
greenwood = CustomMaterial(
texture="WoodGreen",
tex_name="greenwood",
mat_name="greenwood_mat",
tex_attrib=tex_attrib,
mat_attrib=mat_attrib,
)
self.cubeA = BoxObject(
name="cubeA",
size_min=[0.02, 0.02, 0.02],
size_max=[0.02, 0.02, 0.02],
rgba=[1, 0, 0, 1],
material=redwood,
)
self.cubeB = BoxObject(
name="cubeB",
size_min=[0.025, 0.025, 0.025],
size_max=[0.025, 0.025, 0.025],
rgba=[0, 1, 0, 1],
material=greenwood,
)
cubes = [self.cubeA, self.cubeB]
# Create placement initializer
if self.placement_initializer is not None:
self.placement_initializer.reset()
self.placement_initializer.add_objects(cubes)
else:
self.placement_initializer = UniformRandomSampler(
name="ObjectSampler",
mujoco_objects=cubes,
x_range=[-0.08, 0.08],
y_range=[-0.08, 0.08],
rotation=None,
ensure_object_boundary_in_range=False,
ensure_valid_placement=True,
reference_pos=self.table_offset,
z_offset=0.01,
)
# 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=cubes,
)
def _setup_references(self):
"""
Sets up references to important components. A reference is typically an
index or a list of indices that point to the corresponding elements
in a flatten array, which is how MuJoCo stores physical simulation data.
"""
super()._setup_references()
# Additional object references from this env
self.cubeA_body_id = self.sim.model.body_name2id(self.cubeA.root_body)
self.cubeB_body_id = self.sim.model.body_name2id(self.cubeB.root_body)
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():
self.sim.data.set_joint_qpos(
obj.joints[0],
np.concatenate([np.array(obj_pos),
np.array(obj_quat)]))
self.lift_times = 0
self.hold_times = 0
self.align_reward = 0
self.stack_times = 0
def _setup_observables(self):
"""
Sets up observables to be used for this environment. Creates object-based observables if enabled
Returns:
OrderedDict: Dictionary mapping observable names to its corresponding Observable object
"""
observables = super()._setup_observables()
# low-level object information
if self.use_object_obs:
# Get robot prefix and define observables modality
pf = self.robots[0].robot_model.naming_prefix
modality = "object"
# position and rotation of the first cube
@sensor(modality=modality)
def cubeA_pos(obs_cache):
return np.array(self.sim.data.body_xpos[self.cubeA_body_id])
@sensor(modality=modality)
def cubeA_quat(obs_cache):
return convert_quat(np.array(
self.sim.data.body_xquat[self.cubeA_body_id]),
to="xyzw")
@sensor(modality=modality)
def cubeB_pos(obs_cache):
return np.array(self.sim.data.body_xpos[self.cubeB_body_id])
@sensor(modality=modality)
def cubeB_quat(obs_cache):
return convert_quat(np.array(
self.sim.data.body_xquat[self.cubeB_body_id]),
to="xyzw")
@sensor(modality=modality)
def gripper_to_cubeA(obs_cache):
return obs_cache["cubeA_pos"] - obs_cache[f"{pf}eef_pos"] if \
"cubeA_pos" in obs_cache and f"{pf}eef_pos" in obs_cache else np.zeros(3)
@sensor(modality=modality)
def gripper_to_cubeB(obs_cache):
return obs_cache["cubeB_pos"] - obs_cache[f"{pf}eef_pos"] if \
"cubeB_pos" in obs_cache and f"{pf}eef_pos" in obs_cache else np.zeros(3)
@sensor(modality=modality)
def cubeA_to_cubeB(obs_cache):
return obs_cache["cubeB_pos"] - obs_cache["cubeA_pos"] if \
"cubeA_pos" in obs_cache and "cubeB_pos" in obs_cache else np.zeros(3)
sensors = [
cubeA_pos, cubeA_quat, cubeB_pos, cubeB_quat, gripper_to_cubeA,
gripper_to_cubeB, cubeA_to_cubeB
]
names = [s.__name__ for s in sensors]
# Create observables
for name, s in zip(names, sensors):
observables[name] = Observable(
name=name,
sensor=s,
sampling_rate=self.control_freq,
)
return observables
def _check_success(self):
"""
Check if blocks are stacked correctly.
Returns:
bool: True if blocks are correctly stacked
"""
_, _, r_stack = self.staged_rewards()
return r_stack > 0
def visualize(self, vis_settings):
"""
In addition to super call, visualize gripper site proportional to the distance to the cube.
Args:
vis_settings (dict): Visualization keywords mapped to T/F, determining whether that specific
component should be visualized. Should have "grippers" keyword as well as any other relevant
options specified.
"""
# Run superclass method first
super().visualize(vis_settings=vis_settings)
# Color the gripper visualization site according to its distance to the cube
if vis_settings["grippers"]:
self._visualize_gripper_to_target(gripper=self.robots[0].gripper,
target=self.cubeA)
def _post_action(self, action):
"""
In addition to super method, add additional info if requested
Args:
action (np.array): Action to execute within the environment
Returns:
3-tuple:
- (float) reward from the environment
- (bool) whether the current episode is completed or not
- (dict) info about current env step
"""
reward, done, info = super()._post_action(action)
# Update force bias
if self.get_info:
info['add_vals'] = [
'lift_times', 'hold_times', 'align_reward', 'stack_times'
]
info['lift_times'] = self.lift_times
info['hold_times'] = self.hold_times
info['align_reward'] = self.align_reward
info['stack_times'] = self.stack_times
# allow episode to finish early if allowed
return reward, done, info
class LiftNBVCube(SingleArmEnv):
"""
This class corresponds to the lifting task for a single robot arm.
Args:
robots (str or list of str): Specification for specific robot arm(s) to be instantiated within this env
(e.g: "Sawyer" would generate one arm; ["Panda", "Panda", "Sawyer"] would generate three robot arms)
Note: Must be a single single-arm robot!
env_configuration (str): Specifies how to position the robots within the environment (default is "default").
For most single arm environments, this argument has no impact on the robot setup.
controller_configs (str or list of dict): If set, contains relevant controller parameters for creating a
custom controller. Else, uses the default controller for this specific task. Should either be single
dict if same controller is to be used for all robots or else it should be a list of the same length as
"robots" param
gripper_types (str or list of str): type of gripper, used to instantiate
gripper models from gripper factory. Default is "default", which is the default grippers(s) associated
with the robot(s) the 'robots' specification. None removes the gripper, and any other (valid) model
overrides the default gripper. Should either be single str if same gripper type is to be used for all
robots or else it should be a list of the same length as "robots" param
initialization_noise (dict or list of dict): Dict containing the initialization noise parameters.
The expected keys and corresponding value types are specified below:
:`'magnitude'`: The scale factor of uni-variate random noise applied to each of a robot's given initial
joint positions. Setting this value to `None` or 0.0 results in no noise being applied.
If "gaussian" type of noise is applied then this magnitude scales the standard deviation applied,
If "uniform" type of noise is applied then this magnitude sets the bounds of the sampling range
:`'type'`: Type of noise to apply. Can either specify "gaussian" or "uniform"
Should either be single dict if same noise value is to be used for all robots or else it should be a
list of the same length as "robots" param
:Note: Specifying "default" will automatically use the default noise settings.
Specifying None will automatically create the required dict with "magnitude" set to 0.0.
table_full_size (3-tuple): x, y, and z dimensions of the table.
table_friction (3-tuple): the three mujoco friction parameters for
the table.
use_camera_obs (bool): if True, every observation includes rendered image(s)
use_object_obs (bool): if True, include object (cube) information in
the observation.
reward_scale (None or float): Scales the normalized reward function by the amount specified.
If None, environment reward remains unnormalized
reward_shaping (bool): if True, use dense rewards.
placement_initializer (ObjectPositionSampler): if provided, will
be used to place objects on every reset, else a UniformRandomSampler
is used by default.
has_renderer (bool): If true, render the simulation state in
a viewer instead of headless mode.
has_offscreen_renderer (bool): True if using off-screen rendering
render_camera (str): Name of camera to render if `has_renderer` is True. Setting this value to 'None'
will result in the default angle being applied, which is useful as it can be dragged / panned by
the user using the mouse
render_collision_mesh (bool): True if rendering collision meshes in camera. False otherwise.
render_visual_mesh (bool): True if rendering visual meshes in camera. False otherwise.
render_gpu_device_id (int): corresponds to the GPU device id to use for offscreen rendering.
Defaults to -1, in which case the device will be inferred from environment variables
(GPUS or CUDA_VISIBLE_DEVICES).
control_freq (float): how many control signals to receive in every second. This sets the amount of
simulation time that passes between every action input.
horizon (int): Every episode lasts for exactly @horizon timesteps.
ignore_done (bool): True if never terminating the environment (ignore @horizon).
hard_reset (bool): If True, re-loads model, sim, and render object upon a reset call, else,
only calls sim.reset and resets all robosuite-internal variables
camera_names (str or list of str): name of camera to be rendered. Should either be single str if
same name is to be used for all cameras' rendering or else it should be a list of cameras to render.
:Note: At least one camera must be specified if @use_camera_obs is True.
:Note: To render all robots' cameras of a certain type (e.g.: "robotview" or "eye_in_hand"), use the
convention "all-{name}" (e.g.: "all-robotview") to automatically render all camera images from each
robot's camera list).
camera_heights (int or list of int): height of camera frame. Should either be single int if
same height is to be used for all cameras' frames or else it should be a list of the same length as
"camera names" param.
camera_widths (int or list of int): width of camera frame. Should either be single int if
same width is to be used for all cameras' frames or else it should be a list of the same length as
"camera names" param.
camera_depths (bool or list of bool): True if rendering RGB-D, and RGB otherwise. Should either be single
bool if same depth setting is to be used for all cameras or else it should be a list of the same length as
"camera names" param.
Raises:
AssertionError: [Invalid number of robots specified]
"""
def __init__(
self,
robots,
env_configuration="default",
controller_configs=None,
gripper_types="default",
initialization_noise="default",
table_full_size=(0.8, 0.8, 0.05),
table_friction=(1., 5e-3, 1e-4),
use_camera_obs=True,
use_object_obs=True,
reward_scale=1.0,
reward_shaping=False,
placement_initializer=None,
has_renderer=False,
has_offscreen_renderer=True,
render_camera="frontview",
render_collision_mesh=False,
render_visual_mesh=True,
render_gpu_device_id=-1,
control_freq=20,
horizon=1000,
ignore_done=False,
hard_reset=True,
camera_names="agentview",
camera_heights=256,
camera_widths=256,
camera_depths=False,
):
# settings for table top
self.table_full_size = table_full_size
self.table_friction = table_friction
self.table_offset = np.array((0, 0, 0.8))
# reward configuration
self.reward_scale = reward_scale
self.reward_shaping = reward_shaping
# whether to use ground-truth object states
self.use_object_obs = use_object_obs
# object placement initializer
self.placement_initializer = placement_initializer
super().__init__(
robots=robots,
env_configuration=env_configuration,
controller_configs=controller_configs,
mount_types="default",
gripper_types=gripper_types,
initialization_noise=initialization_noise,
use_camera_obs=use_camera_obs,
has_renderer=has_renderer,
has_offscreen_renderer=has_offscreen_renderer,
render_camera=render_camera,
render_collision_mesh=render_collision_mesh,
render_visual_mesh=render_visual_mesh,
render_gpu_device_id=render_gpu_device_id,
control_freq=control_freq,
horizon=horizon,
ignore_done=ignore_done,
hard_reset=hard_reset,
camera_names=camera_names,
camera_heights=camera_heights,
camera_widths=camera_widths,
camera_depths=camera_depths,
)
def reward(self, action=None):
"""
Reward function for the task.
Sparse un-normalized reward:
- a discrete reward of 2.25 is provided if the cube is lifted
Un-normalized summed components if using reward shaping:
- Reaching: in [0, 1], to encourage the arm to reach the cube
- Grasping: in {0, 0.25}, non-zero if arm is grasping the cube
- Lifting: in {0, 1}, non-zero if arm has lifted the cube
The sparse reward only consists of the lifting component.
Note that the final reward is normalized and scaled by
reward_scale / 2.25 as well so that the max score is equal to reward_scale
Args:
action (np array): [NOT USED]
Returns:
float: reward value
"""
reward = 0.
# sparse completion reward
if self._check_success():
reward = 2.25
# use a shaping reward
elif self.reward_shaping:
# reaching reward
cube_pos = self.sim.data.body_xpos[self.cube_body_id]
gripper_site_pos = self.sim.data.site_xpos[self.robots[0].eef_site_id]
dist = np.linalg.norm(gripper_site_pos - cube_pos)
reaching_reward = 1 - np.tanh(10.0 * dist)
reward += reaching_reward
# grasping reward
if self._check_grasp(gripper=self.robots[0].gripper, object_geoms=self.cube):
reward += 0.25
# Scale reward if requested
if self.reward_scale is not None:
reward *= self.reward_scale / 2.25
return reward
def _load_model(self):
"""
Loads an xml model, puts it in self.model
"""
super()._load_model()
# Adjust base pose accordingly
xpos = self.robots[0].robot_model.base_xpos_offset["table"](self.table_full_size[0])
self.robots[0].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
tex_attrib = {
"type": "cube",
}
mat_attrib = {
"texrepeat": "1 1",
"specular": "0.4",
"shininess": "0.1",
}
redwood = CustomMaterial(
texture="WoodRed",
tex_name="redwood",
mat_name="redwood_mat",
tex_attrib=tex_attrib,
mat_attrib=mat_attrib,
)
#self.cube = BoxObject(
# name="cube",
# size_min=[0.020, 0.020, 0.020], # [0.015, 0.015, 0.015],
# size_max=[0.022, 0.022, 0.022], # [0.018, 0.018, 0.018])
# rgba=[1, 0, 0, 1],
# material=redwood,
#)
self.cube = NBVCubeObject(name="cube")
# Create placement initializer
if self.placement_initializer is not None:
self.placement_initializer.reset()
self.placement_initializer.add_objects(self.cube)
else:
self.placement_initializer = UniformRandomSampler(
name="ObjectSampler",
mujoco_objects=self.cube,
x_range=[-0.03, 0.03],
y_range=[-0.03, 0.03],
rotation=None,
ensure_object_boundary_in_range=False,
ensure_valid_placement=True,
reference_pos=self.table_offset,
z_offset=0.01,
)
# 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.cube,
)
def _get_reference(self):
"""
Sets up references to important components. A reference is typically an
index or a list of indices that point to the corresponding elements
in a flatten array, which is how MuJoCo stores physical simulation data.
"""
super()._get_reference()
# Additional object references from this env
self.cube_body_id = self.sim.model.body_name2id(self.cube.root_body)
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():
self.sim.data.set_joint_qpos(obj.joints[0], np.concatenate([np.array(obj_pos), np.array(obj_quat)]))
def _get_observation(self):
"""
Returns an OrderedDict containing observations [(name_string, np.array), ...].
Important keys:
`'robot-state'`: contains robot-centric information.
`'object-state'`: requires @self.use_object_obs to be True. Contains object-centric information.
`'image'`: requires @self.use_camera_obs to be True. Contains a rendered frame from the simulation.
`'depth'`: requires @self.use_camera_obs and @self.camera_depth to be True.
Contains a rendered depth map from the simulation
Returns:
OrderedDict: Observations from the environment
"""
di = super()._get_observation()
# low-level object information
if self.use_object_obs:
# Get robot prefix
pr = self.robots[0].robot_model.naming_prefix
# position and rotation of object
cube_pos = np.array(self.sim.data.body_xpos[self.cube_body_id])
cube_quat = convert_quat(
np.array(self.sim.data.body_xquat[self.cube_body_id]), to="xyzw"
)
di["cube_pos"] = cube_pos
di["cube_quat"] = cube_quat
gripper_site_pos = np.array(self.sim.data.site_xpos[self.robots[0].eef_site_id])
di[pr + "gripper_to_cube"] = gripper_site_pos - cube_pos
di["object-state"] = np.concatenate(
[cube_pos, cube_quat, di[pr + "gripper_to_cube"]]
)
return di
def visualize(self, vis_settings):
"""
In addition to super call, visualize gripper site proportional to the distance to the cube.
Args:
vis_settings (dict): Visualization keywords mapped to T/F, determining whether that specific
component should be visualized. Should have "grippers" keyword as well as any other relevant
options specified.
"""
# Run superclass method first
super().visualize(vis_settings=vis_settings)
# Color the gripper visualization site according to its distance to the cube
if vis_settings["grippers"]:
self._visualize_gripper_to_target(gripper=self.robots[0].gripper, target=self.cube)
def _check_success(self):
"""
Check if cube has been lifted.
Returns:
bool: True if cube has been lifted
"""
cube_height = self.sim.data.body_xpos[self.cube_body_id][2]
table_height = self.model.mujoco_arena.table_offset[2]
# cube is higher than the table top above a margin
return cube_height > table_height + 0.04
class Ultrasound(SingleArmEnv):
"""
This class corresponds to the ultrasound task for a single robot arm.
Args:
robots (str or list of str): Specification for specific robot arm(s) to be instantiated within this env
(e.g: "Sawyer" would generate one arm; ["Panda", "Panda", "Sawyer"] would generate three robot arms)
Note: Must be a single single-arm robot!
env_configuration (str): Specifies how to position the robots within the environment (default is "default").
For most single arm environments, this argument has no impact on the robot setup.
controller_configs (str or list of dict): If set, contains relevant controller parameters for creating a
custom controller. Else, uses the default controller for this specific task. Should either be single
dict if same controller is to be used for all robots or else it should be a list of the same length as
"robots" param
gripper_types (str or list of str): type of gripper, used to instantiate
gripper models from gripper factory. Default is "default", which is the default grippers(s) associated
with the robot(s) the 'robots' specification. None removes the gripper, and any other (valid) model
overrides the default gripper. Should either be single str if same gripper type is to be used for all
robots or else it should be a list of the same length as "robots" param
table_full_size (3-tuple): x, y, and z dimensions of the table.
table_friction (3-tuple): the three mujoco friction parameters for
the table.
use_camera_obs (bool): if True, every observation includes rendered image(s)
use_object_obs (bool): if True, include object (cube) information in
the observation.
placement_initializer (ObjectPositionSampler): if provided, will
be used to place objects on every reset, else a UniformRandomSampler
is used by default.
has_renderer (bool): If true, render the simulation state in
a viewer instead of headless mode.
has_offscreen_renderer (bool): True if using off-screen rendering
render_camera (str): Name of camera to render if `has_renderer` is True. Setting this value to 'None'
will result in the default angle being applied, which is useful as it can be dragged / panned by
the user using the mouse
render_collision_mesh (bool): True if rendering collision meshes in camera. False otherwise.
render_visual_mesh (bool): True if rendering visual meshes in camera. False otherwise.
render_gpu_device_id (int): corresponds to the GPU device id to use for offscreen rendering.
Defaults to -1, in which case the device will be inferred from environment variables
(GPUS or CUDA_VISIBLE_DEVICES).
control_freq (float): how many control signals to receive in every second. This sets the amount of
simulation time that passes between every action input.
horizon (int): Every episode lasts for exactly @horizon timesteps.
ignore_done (bool): True if never terminating the environment (ignore @horizon).
hard_reset (bool): If True, re-loads model, sim, and render object upon a reset call, else,
only calls sim.reset and resets all robosuite-internal variables
camera_names (str or list of str): name of camera to be rendered. Should either be single str if
same name is to be used for all cameras' rendering or else it should be a list of cameras to render.
:Note: At least one camera must be specified if @use_camera_obs is True.
:Note: To render all robots' cameras of a certain type (e.g.: "robotview" or "eye_in_hand"), use the
convention "all-{name}" (e.g.: "all-robotview") to automatically render all camera images from each
robot's camera list).
camera_heights (int or list of int): height of camera frame. Should either be single int if
same height is to be used for all cameras' frames or else it should be a list of the same length as
"camera names" param.
camera_widths (int or list of int): width of camera frame. Should either be single int if
same width is to be used for all cameras' frames or else it should be a list of the same length as
"camera names" param.
camera_depths (bool or list of bool): True if rendering RGB-D, and RGB otherwise. Should either be single
bool if same depth setting is to be used for all cameras or else it should be a list of the same length as
"camera names" param.
early_termination (bool): If True, episode is allowed to finish early.
save_data (bool): If True, data from the episode is collected and saved.
deterministic_trajectory (bool): If True, chooses a deterministic trajectory which goes along the x-axis of the torso.
torso_solref_randomization (bool): If True, randomize the stiffness and damping parameter of the torso.
initial_probe_pos_randomization (bool): If True, Gaussian noise will be added to the initial position of the probe.
use_box_torso (bool): If True, use a box shaped soft body. Else, use a cylinder shaped soft body.
Raises:
AssertionError: [Invalid number of robots specified]
"""
def __init__(
self,
robots,
env_configuration="default",
controller_configs=None,
gripper_types="UltrasoundProbeGripper",
initialization_noise="default",
table_full_size=(0.8, 0.8, 0.05),
table_friction=100 * (1., 5e-3, 1e-4),
use_camera_obs=True,
use_object_obs=True,
reward_scale=1.0,
reward_shaping=False,
placement_initializer=None,
has_renderer=False,
has_offscreen_renderer=True,
render_camera="frontview",
render_collision_mesh=False,
render_visual_mesh=True,
render_gpu_device_id=-1,
control_freq=20,
horizon=1000,
ignore_done=False,
hard_reset=True,
camera_names="agentview",
camera_heights=256,
camera_widths=256,
camera_depths=False,
early_termination=False,
save_data=False,
deterministic_trajectory=False,
torso_solref_randomization=False,
initial_probe_pos_randomization=False,
use_box_torso=True,
):
assert gripper_types == "UltrasoundProbeGripper",\
"Tried to specify gripper other than UltrasoundProbeGripper in Ultrasound environment!"
assert robots == "UR5e" or robots == "Panda", \
"Robot must be UR5e or Panda!"
assert "OSC" or "HMFC" in controller_configs["type"], \
"The robot controller must be of type OSC or HMFC"
# settings for table top
self.table_full_size = table_full_size
self.table_friction = table_friction
self.table_offset = np.array((0, 0, 0.8))
# settings for joint initialization noise (Gaussian)
self.mu = 0
self.sigma = 0.010
# settings for contact force running mean
self.alpha = 0.1 # decay factor (high alpha -> discounts older observations faster). Must be in (0, 1)
# reward configuration
self.reward_scale = reward_scale
self.reward_shaping = reward_shaping
# error multipliers
self.pos_error_mul = 90
self.ori_error_mul = 0.2
self.vel_error_mul = 45
self.force_error_mul = 0.7
self.der_force_error_mul = 0.01
# reward multipliers
self.pos_reward_mul = 5
self.ori_reward_mul = 1
self.vel_reward_mul = 1
self.force_reward_mul = 3
self.der_force_reward_mul = 2
# desired states
self.goal_quat = np.array([
-0.69192486, 0.72186726, -0.00514253, -0.01100909
]) # Upright probe orientation found from experimenting (x,y,z,w)
self.goal_velocity = 0.04 # norm of velocity vector
self.goal_contact_z_force = 5 # (N)
self.goal_der_contact_z_force = 0 # derivative of contact force
# early termination configuration
self.pos_error_threshold = 1.0
self.ori_error_threshold = 0.10
# examination trajectory
self.top_torso_offset = 0.039 if use_box_torso else 0.041 # offset from z_center of torso to top of torso
self.x_range = 0.15 # how large the torso is from center to end in x-direction
self.y_range = 0.09 if use_box_torso else 0.05 # how large the torso is from center to end in y-direction
self.grid_pts = 50 # how many points in the grid
# whether to use ground-truth object states
self.use_object_obs = use_object_obs
# object placement initializer
self.placement_initializer = placement_initializer
# randomization settings
self.torso_solref_randomization = torso_solref_randomization
self.initial_probe_pos_randomization = initial_probe_pos_randomization
# misc settings
self.early_termination = early_termination
self.save_data = save_data
self.deterministic_trajectory = deterministic_trajectory
self.use_box_torso = use_box_torso
super().__init__(
robots=robots,
env_configuration=env_configuration,
controller_configs=controller_configs,
mount_types="default",
gripper_types=gripper_types,
initialization_noise=None,
use_camera_obs=use_camera_obs,
has_renderer=has_renderer,
has_offscreen_renderer=has_offscreen_renderer,
render_camera=render_camera,
render_collision_mesh=render_collision_mesh,
render_visual_mesh=render_visual_mesh,
render_gpu_device_id=render_gpu_device_id,
control_freq=control_freq,
horizon=horizon,
ignore_done=ignore_done,
hard_reset=hard_reset,
camera_names=camera_names,
camera_heights=camera_heights,
camera_widths=camera_widths,
camera_depths=camera_depths,
)
def reward(self, action=None):
"""
Reward function for the task.
Args:
action (np array): [NOT USED]
Returns:
float: reward value
"""
reward = 0.
ee_current_ori = convert_quat(self._eef_xquat,
to="wxyz") # (w, x, y, z) quaternion
ee_desired_ori = convert_quat(self.goal_quat, to="wxyz")
# position
self.pos_error = np.square(self.pos_error_mul *
(self._eef_xpos[0:-1] - self.traj_pt[0:-1]))
self.pos_reward = self.pos_reward_mul * np.exp(
-1 * np.linalg.norm(self.pos_error))
# orientation
self.ori_error = self.ori_error_mul * distance_quat(
ee_current_ori, ee_desired_ori)
self.ori_reward = self.ori_reward_mul * np.exp(-1 * self.ori_error)
# velocity
self.vel_error = np.square(
self.vel_error_mul * (self.vel_running_mean - self.goal_velocity))
self.vel_reward = self.vel_reward_mul * np.exp(
-1 * np.linalg.norm(self.vel_error))
# force
self.force_error = np.square(
self.force_error_mul *
(self.z_contact_force_running_mean - self.goal_contact_z_force))
self.force_reward = self.force_reward_mul * np.exp(
-1 *
self.force_error) if self._check_probe_contact_with_torso() else 0
# derivative force
self.der_force_error = np.square(
self.der_force_error_mul *
(self.der_z_contact_force - self.goal_der_contact_z_force))
self.der_force_reward = self.der_force_reward_mul * np.exp(
-1 * self.der_force_error) if self._check_probe_contact_with_torso(
) else 0
# add rewards
reward += (self.pos_reward + self.ori_reward + self.vel_reward +
self.force_reward + self.der_force_reward)
return reward
def _load_model(self):
"""
Loads an xml model, puts it in self.model
"""
super()._load_model()
# Adjust base pose accordingly
xpos = self.robots[0].robot_model.base_xpos_offset["table"](
self.table_full_size[0])
self.robots[0].robot_model.set_base_xpos(xpos)
# Load model for table top workspace
mujoco_arena = UltrasoundArena()
# Arena always gets set to zero origin
mujoco_arena.set_origin([0, 0, 0])
# Initialize torso object
self.torso = SoftBoxObject(
name="torso") if self.use_box_torso else SoftTorsoObject(
name="torso")
if self.torso_solref_randomization:
# Randomize torso's stiffness and damping (values are taken from my project thesis)
stiffness = np.random.randint(1300, 1600)
damping = np.random.randint(17, 41)
self.torso.set_damping(damping)
self.torso.set_stiffness(stiffness)
# Create placement initializer
if self.placement_initializer is not None:
self.placement_initializer.reset()
self.placement_initializer.add_objects(self.torso)
else:
self.placement_initializer = UniformRandomSampler(
name="ObjectSampler",
mujoco_objects=[self.torso],
x_range=[0, 0], #[-0.12, 0.12],
y_range=[0, 0], #[-0.12, 0.12],
rotation=None,
ensure_object_boundary_in_range=False,
ensure_valid_placement=True,
reference_pos=self.table_offset,
z_offset=0.005,
)
# task includes arena, robot, and objects of interest
self.model = UltrasoundTask(
mujoco_arena=mujoco_arena,
mujoco_robots=[robot.robot_model for robot in self.robots],
mujoco_objects=[self.torso])
def _setup_references(self):
"""
Sets up references to important components. A reference is typically an
index or a list of indices that point to the corresponding elements
in a flatten array, which is how MuJoCo stores physical simulation data.
"""
super()._setup_references()
# additional object references from this env
self.torso_body_id = self.sim.model.body_name2id(self.torso.root_body)
self.probe_id = self.sim.model.body_name2id(
self.robots[0].gripper.root_body)
def _setup_observables(self):
"""
Sets up observables to be used for this environment. Creates object-based observables if enabled
Returns:
OrderedDict: Dictionary mapping observable names to its corresponding Observable object
"""
observables = super()._setup_observables()
pf = self.robots[0].robot_model.naming_prefix
# Remove unnecessary observables
del observables[pf + "joint_pos"]
del observables[pf + "joint_pos_cos"]
del observables[pf + "joint_pos_sin"]
del observables[pf + "joint_vel"]
del observables[pf + "gripper_qvel"]
del observables[pf + "gripper_qpos"]
del observables[pf + "eef_pos"]
del observables[pf + "eef_quat"]
sensors = []
# probe information
modality = f"{pf}proprio" # Need to use this modality since proprio obs cannot be empty in GymWrapper
@sensor(modality=modality)
def eef_contact_force(obs_cache):
return self.sim.data.cfrc_ext[self.probe_id][-3:]
@sensor(modality=modality)
def eef_torque(obs_cache):
return self.robots[0].ee_torque
@sensor(modality=modality)
def eef_vel(obs_cache):
return self.robots[0]._hand_vel
@sensor(modality=modality)
def eef_contact_force_z_diff(obs_cache):
return self.z_contact_force_running_mean - self.goal_contact_z_force
@sensor(modality=modality)
def eef_contact_derivative_force_z_diff(obs_cache):
return self.der_z_contact_force - self.goal_der_contact_z_force
@sensor(modality=modality)
def eef_vel_diff(obs_cache):
return self.vel_running_mean - self.goal_velocity
@sensor(modality=modality)
def eef_pose_diff(obs_cache):
pos_error = self._eef_xpos - self.traj_pt
quat_error = difference_quat(self._eef_xquat, self.goal_quat)
pose_error = np.concatenate((pos_error, quat_error))
return pose_error
sensors += [
eef_contact_force, eef_torque, eef_vel, eef_contact_force_z_diff,
eef_contact_derivative_force_z_diff, eef_vel_diff, eef_pose_diff
]
names = [s.__name__ for s in sensors]
# Create observables
for name, s in zip(names, sensors):
observables[name] = Observable(
name=name,
sensor=s,
sampling_rate=self.control_freq,
)
return observables
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 in object_placements.values():
self.sim.data.set_joint_qpos(
obj.joints[0],
np.concatenate(
[np.array(obj_pos),
np.array([0.5, 0.5, -0.5, -0.5])]))
self.sim.forward() # update sim states
# says if probe has been in touch with torso
self.has_touched_torso = False
# initial position of end-effector
self.ee_initial_pos = self._eef_xpos
# create trajectory
self.trajectory = self.get_trajectory()
# initialize trajectory step
self.initial_traj_step = np.random.default_rng().uniform(
low=0, high=self.num_waypoints - 1)
self.traj_step = self.initial_traj_step # step at which to evaluate trajectory. Must be in interval [0, num_waypoints - 1]
# set first trajectory point
self.traj_pt = self.trajectory.eval(self.traj_step)
self.traj_pt_vel = self.trajectory.deriv(self.traj_step)
# give controller access to robot (and its measurements)
if self.robots[0].controller.name == "HMFC":
self.robots[0].controller.set_robot(self.robots[0])
# initialize controller's trajectory
self.robots[0].controller.traj_pos = self.traj_pt
self.robots[0].controller.traj_ori = T.quat2axisangle(self.goal_quat)
# get initial joint positions for robot
init_qpos = self._get_initial_qpos()
# override initial robot joint positions
self.robots[0].set_robot_joint_positions(init_qpos)
# update controller with new initial joints
self.robots[0].controller.update_initial_joints(init_qpos)
# initialize previously contact force measurement
self.prev_z_contact_force = 0
# intialize derivative of contact force
self.der_z_contact_force = 0
# initialize running mean of velocity
self.vel_running_mean = np.linalg.norm(self.robots[0]._hand_vel)
# initialize running mean of contact force
self.z_contact_force_running_mean = self.sim.data.cfrc_ext[
self.probe_id][-1]
# initialize data collection
if self.save_data:
# simulation data
self.data_ee_pos = np.array(np.zeros((self.horizon, 3)))
self.data_ee_goal_pos = np.array(np.zeros((self.horizon, 3)))
self.data_ee_ori_diff = np.array(np.zeros(self.horizon))
self.data_ee_vel = np.array(np.zeros((self.horizon, 3)))
self.data_ee_goal_vel = np.array(np.zeros(self.horizon))
self.data_ee_running_mean_vel = np.array(np.zeros(self.horizon))
self.data_ee_quat = np.array(np.zeros(
(self.horizon, 4))) # (x,y,z,w)
self.data_ee_goal_quat = np.array(np.zeros(
(self.horizon, 4))) # (x,y,z,w)
self.data_ee_diff_quat = np.array(np.zeros(
self.horizon)) # (x,y,z,w)
self.data_ee_z_contact_force = np.array(np.zeros(self.horizon))
self.data_ee_z_goal_contact_force = np.array(np.zeros(
self.horizon))
self.data_ee_z_running_mean_contact_force = np.array(
np.zeros(self.horizon))
self.data_ee_z_derivative_contact_force = np.array(
np.zeros(self.horizon))
self.data_ee_z_goal_derivative_contact_force = np.array(
np.zeros(self.horizon))
self.data_is_contact = np.array(np.zeros(self.horizon))
self.data_q_pos = np.array(
np.zeros((self.horizon, self.robots[0].dof)))
self.data_q_torques = np.array(
np.zeros((self.horizon, self.robots[0].dof)))
self.data_time = np.array(np.zeros(self.horizon))
# reward data
self.data_pos_reward = np.array(np.zeros(self.horizon))
self.data_ori_reward = np.array(np.zeros(self.horizon))
self.data_vel_reward = np.array(np.zeros(self.horizon))
self.data_force_reward = np.array(np.zeros(self.horizon))
self.data_der_force_reward = np.array(np.zeros(self.horizon))
# policy/controller data
self.data_action = np.array(
np.zeros((self.horizon, self.robots[0].action_dim)))
def _post_action(self, action):
"""
In addition to super method, add additional info if requested
Args:
action (np.array): Action to execute within the environment
Returns:
3-tuple:
- (float) reward from the environment
- (bool) whether the current episode is completed or not
- (dict) info about current env step
"""
reward, done, info = super()._post_action(action)
# Convert to trajectory timstep
normalizer = (self.horizon / (self.num_waypoints - 1)
) # equally many timesteps to reach each waypoint
self.traj_step = self.timestep / normalizer + self.initial_traj_step
# update trajectory point
self.traj_pt = self.trajectory.eval(self.traj_step)
# update controller's trajectory
self.robots[0].controller.traj_pos = self.traj_pt
# update velocity running mean (simple moving average)
self.vel_running_mean += ((np.linalg.norm(self.robots[0]._hand_vel) -
self.vel_running_mean) / self.timestep)
# update derivative of contact force
z_contact_force = self.sim.data.cfrc_ext[self.probe_id][-1]
self.der_z_contact_force = (z_contact_force - self.prev_z_contact_force
) / self.control_timestep
self.prev_z_contact_force = z_contact_force
# update contact force running mean (exponential moving average)
self.z_contact_force_running_mean = self.alpha * z_contact_force + (
1 - self.alpha) * self.z_contact_force_running_mean
# check for early termination
if self.early_termination:
done = done or self._check_terminated()
# collect data
if self.save_data:
# simulation data
self.data_ee_pos[self.timestep - 1] = self._eef_xpos
self.data_ee_goal_pos[self.timestep - 1] = self.traj_pt
self.data_ee_vel[self.timestep - 1] = self.robots[0]._hand_vel
self.data_ee_goal_vel[self.timestep - 1] = self.goal_velocity
self.data_ee_running_mean_vel[self.timestep -
1] = self.vel_running_mean
self.data_ee_quat[self.timestep - 1] = self._eef_xquat
self.data_ee_goal_quat[self.timestep - 1] = self.goal_quat
self.data_ee_diff_quat[self.timestep - 1] = distance_quat(
convert_quat(self._eef_xquat, to="wxyz"),
convert_quat(self.goal_quat, to="wxyz"))
self.data_ee_z_contact_force[
self.timestep - 1] = self.sim.data.cfrc_ext[self.probe_id][-1]
self.data_ee_z_goal_contact_force[self.timestep -
1] = self.goal_contact_z_force
self.data_ee_z_running_mean_contact_force[
self.timestep - 1] = self.z_contact_force_running_mean
self.data_ee_z_derivative_contact_force[
self.timestep - 1] = self.der_z_contact_force
self.data_ee_z_goal_derivative_contact_force[
self.timestep - 1] = self.goal_der_contact_z_force
self.data_is_contact[self.timestep -
1] = self._check_probe_contact_with_torso()
self.data_q_pos[self.timestep -
1] = self.robots[0]._joint_positions
self.data_q_torques[self.timestep - 1] = self.robots[0].torques
self.data_time[self.timestep - 1] = (
self.timestep -
1) / self.horizon * 100 # percentage of completed episode
# reward data
self.data_pos_reward[self.timestep - 1] = self.pos_reward
self.data_ori_reward[self.timestep - 1] = self.ori_reward
self.data_vel_reward[self.timestep - 1] = self.vel_reward
self.data_force_reward[self.timestep - 1] = self.force_reward
self.data_der_force_reward[self.timestep -
1] = self.der_force_reward
# policy/controller data
self.data_action[self.timestep - 1] = action
# save data
if done and self.save_data:
# simulation data
sim_data_fldr = "simulation_data"
self._save_data(self.data_ee_pos, sim_data_fldr, "ee_pos")
self._save_data(self.data_ee_goal_pos, sim_data_fldr,
"ee_goal_pos")
self._save_data(self.data_ee_vel, sim_data_fldr, "ee_vel")
self._save_data(self.data_ee_goal_vel, sim_data_fldr,
"ee_goal_vel")
self._save_data(self.data_ee_running_mean_vel, sim_data_fldr,
"ee_running_mean_vel")
self._save_data(self.data_ee_quat, sim_data_fldr, "ee_quat")
self._save_data(self.data_ee_goal_quat, sim_data_fldr,
"ee_goal_quat")
self._save_data(self.data_ee_diff_quat, sim_data_fldr,
"ee_diff_quat")
self._save_data(self.data_ee_z_contact_force, sim_data_fldr,
"ee_z_contact_force")
self._save_data(self.data_ee_z_goal_contact_force, sim_data_fldr,
"ee_z_goal_contact_force")
self._save_data(self.data_ee_z_running_mean_contact_force,
sim_data_fldr, "ee_z_running_mean_contact_force")
self._save_data(self.data_ee_z_derivative_contact_force,
sim_data_fldr, "ee_z_derivative_contact_force")
self._save_data(self.data_ee_z_goal_derivative_contact_force,
sim_data_fldr,
"ee_z_goal_derivative_contact_force")
self._save_data(self.data_is_contact, sim_data_fldr, "is_contact")
self._save_data(self.data_q_pos, sim_data_fldr, "q_pos")
self._save_data(self.data_q_torques, sim_data_fldr, "q_torques")
self._save_data(self.data_time, sim_data_fldr, "time")
# reward data
reward_data_fdlr = "reward_data"
self._save_data(self.data_pos_reward, reward_data_fdlr, "pos")
self._save_data(self.data_ori_reward, reward_data_fdlr, "ori")
self._save_data(self.data_vel_reward, reward_data_fdlr, "vel")
self._save_data(self.data_force_reward, reward_data_fdlr, "force")
self._save_data(self.data_der_force_reward, reward_data_fdlr,
"derivative_force")
# policy/controller data
self._save_data(self.data_action, "policy_data", "action")
return reward, done, info
def visualize(self, vis_settings):
"""
Args:
vis_settings (dict): Visualization keywords mapped to T/F, determining whether that specific
component should be visualized. Should have "grippers" keyword as well as any other relevant
options specified.
"""
# Run superclass method first
super().visualize(vis_settings=vis_settings)
def _check_success(self):
return False
def _check_terminated(self):
"""
Check if the task has completed one way or another. The following conditions lead to termination:
- Collision with table
- Joint Limit reached
- Deviates from trajectory position
- Deviates from desired orientation when in contact with torso
- Loses contact with torso
Returns:
bool: True if episode is terminated
"""
terminated = False
# Prematurely terminate if reaching joint limits
if self.robots[0].check_q_limits():
print(40 * '-' + " JOINT LIMIT " + 40 * '-')
terminated = True
# Prematurely terminate if probe deviates away from trajectory (represented by a low position reward)
if np.linalg.norm(self.pos_error) > self.pos_error_threshold:
print(40 * '-' + " DEVIATES FROM TRAJECTORY " + 40 * '-')
terminated = True
# Prematurely terminate if probe deviates from desired orientation when touching probe
if self._check_probe_contact_with_torso(
) and self.ori_error > self.ori_error_threshold:
print(40 * '-' +
" (TOUCHING BODY) PROBE DEVIATES FROM DESIRED ORIENTATION " +
40 * '-')
terminated = True
# Prematurely terminate if probe loses contact with torso
if self.has_touched_torso and not self._check_probe_contact_with_torso(
):
print(40 * '-' + " LOST CONTACT WITH TORSO " + 40 * '-')
terminated = True
return terminated
def _get_contacts_objects(self, model):
"""
Checks for any contacts with @model (as defined by @model's contact_geoms) and returns the set of
contact objects currently in contact with that model (excluding the geoms that are part of the model itself).
Args:
model (MujocoModel): Model to check contacts for.
Returns:
set: Unique contact objects containing information about contacts with this model.
Raises:
AssertionError: [Invalid input type]
"""
# Make sure model is MujocoModel type
assert isinstance(model, MujocoModel), \
"Inputted model must be of type MujocoModel; got type {} instead!".format(type(model))
contact_set = set()
for contact in self.sim.data.contact[:self.sim.data.ncon]:
# check contact geom in geoms; add to contact set if match is found
g1, g2 = self.sim.model.geom_id2name(
contact.geom1), self.sim.model.geom_id2name(contact.geom2)
if g1 in model.contact_geoms or g2 in model.contact_geoms:
contact_set.add(contact)
return contact_set
def _check_probe_contact_with_upper_part_torso(self):
"""
Check if the probe is in contact with the upper/top part of torso. Touching the torso on the sides should not count as contact.
Returns:
bool: True if probe both is in contact with upper part of torso and inside distance threshold from the torso center.
"""
# check for contact only if probe is in contact with upper part and close to torso center
if self._eef_xpos[-1] >= self._torso_xpos[-1] and np.linalg.norm(
self._eef_xpos[:2] - self._torso_xpos[:2]) < 0.14:
return self._check_probe_contact_with_torso()
return False
def _check_probe_contact_with_torso(self):
"""
Check if the probe is in contact with the torso.
NOTE This method utilizes the autogenerated geom names for MuJoCo-native composite objects
Returns:
bool: True if probe is in contact with torso
"""
gripper_contacts = self._get_contacts_objects(self.robots[0].gripper)
reg_ex = "[G]\d+[_]\d+[_]\d+$"
# check contact with torso geoms based on autogenerated names
for contact in gripper_contacts:
g1, g2 = self.sim.model.geom_id2name(
contact.geom1), self.sim.model.geom_id2name(contact.geom2)
match1 = re.search(reg_ex, g1)
match2 = re.search(reg_ex, g2)
if match1 != None or match2 != None:
contact_normal_axis = contact.frame[:3]
self.has_touched_torso = True
return True
return False
def _check_probe_contact_with_table(self):
"""
Check if the probe is in contact with the tabletop.
Returns:
bool: True if probe is in contact with table
"""
return self.check_contact(self.robots[0].gripper, "table_collision")
def get_trajectory(self):
"""
Calculates a trajectory between two waypoints on the torso. The waypoints are extracted from a grid on the torso.
The first waypoint is given at time t=0, and the second waypoint is given at t=1.
If self.deterministic_trajectory is true, a deterministic trajectory along the x-axis of the torso is calculated.
Args:
Returns:
(klampt.model.trajectory Object): trajectory
"""
grid = self._get_torso_grid()
if self.deterministic_trajectory:
start_point = [0.062, -0.020, 0.896]
end_point = [-0.032, -0.075, 0.896]
#start_point = [grid[0, 0], grid[1, 4], self._torso_xpos[-1] + self.top_torso_offset]
#end_point = [grid[0, int(self.grid_pts / 2) - 1], grid[1, 5], self._torso_xpos[-1] + self.top_torso_offset]
else:
start_point = self._get_waypoint(grid)
end_point = self._get_waypoint(grid)
milestones = np.array([start_point, end_point])
self.num_waypoints = np.size(milestones, 0)
return trajectory.Trajectory(milestones=milestones)
def _get_torso_grid(self):
"""
Creates a 2D grid in the xy-plane on the top of the torso.
Args:
Returns:
(numpy.array): grid. First row contains x-coordinates and the second row contains y-coordinates.
"""
x = np.linspace(
-self.x_range + self._torso_xpos[0] + 0.03,
self.x_range + self._torso_xpos[0],
num=self.grid_pts
) # add offset in negative range due to weird robot angles close to robot base
y = np.linspace(-self.y_range + self._torso_xpos[1],
self.y_range + self._torso_xpos[1],
num=self.grid_pts)
x = np.array([x])
y = np.array([y])
return np.concatenate((x, y))
def _get_waypoint(self, grid):
"""
Extracts a random waypoint from the grid.
Args:
Returns:
(numpy.array): waypoint
"""
x_pos = np.random.choice(grid[0])
y_pos = np.random.choice(grid[1])
z_pos = self._torso_xpos[-1] + self.top_torso_offset
return np.array([x_pos, y_pos, z_pos])
def _get_initial_qpos(self):
"""
Calculates the initial joint position for the robot based on the initial desired pose (self.traj_pt, self.goal_quat).
If self.initial_probe_pos_randomization is True, Guassian noise is added to the initial position of the probe.
Args:
Returns:
(np.array): n joint positions
"""
pos = np.array(self.traj_pt)
if self.initial_probe_pos_randomization:
pos = self._add_noise_to_pos(pos)
pos = self._convert_robosuite_to_toolbox_xpos(pos)
ori_euler = mat2euler(quat2mat(self.goal_quat))
# desired pose
T = SE3(pos) * SE3.RPY(ori_euler)
# find initial joint positions
if self.robots[0].name == "UR5e":
robot = rtb.models.DH.UR5()
sol = robot.ikine_min(T, q0=self.robots[0].init_qpos)
# flip last joint around (pi)
sol.q[-1] -= np.pi
return sol.q
elif self.robots[0].name == "Panda":
robot = rtb.models.DH.Panda()
sol = robot.ikine_min(T, q0=self.robots[0].init_qpos)
return sol.q
def _convert_robosuite_to_toolbox_xpos(self, pos):
"""
Converts origin used in robosuite to origin used for robotics toolbox. Also transforms robosuite world frame (vectors x, y, z) to
to correspond to world frame used in toolbox.
Args:
pos (np.array): position (x,y,z) given in robosuite coordinates and frame
Returns:
(np.array): position (x,y,z) given in robotics toolbox coordinates and frame
"""
xpos_offset = self.robots[0].robot_model.base_xpos_offset["table"](
self.table_full_size[0])[0]
zpos_offset = self.robots[0].robot_model.top_offset[-1] - 0.016
# the numeric offset values have been found empirically, where they are chosen so that
# self._eef_xpos matches the desired position.
if self.robots[0].name == "UR5e":
return np.array([
-pos[0] + xpos_offset + 0.08, -pos[1] + 0.025,
pos[2] - zpos_offset + 0.15
])
if self.robots[0].name == "Panda":
return np.array([
pos[0] - xpos_offset - 0.06, pos[1],
pos[2] - zpos_offset + 0.111
])
def _add_noise_to_pos(self, init_pos):
"""
Adds Gaussian noise (variance) to the position.
Args:
init_pos (np.array): initial probe position
Returns:
(np.array): position with added noise
"""
z_noise = np.random.normal(self.mu, self.sigma, 1)
xy_noise = np.random.normal(self.mu, self.sigma / 4, 2)
x = init_pos[0] + xy_noise[0]
y = init_pos[1] + xy_noise[1]
z = init_pos[2] + z_noise[0]
return np.array([x, y, z])
def _save_data(self, data, fldr, filename):
"""
Saves data to desired path.
Args:
data (np.array): Data to be saved
fldr (string): Name of destination folder
filename (string): Name of file
Returns:
"""
os.makedirs(fldr, exist_ok=True)
idx = 1
path = os.path.join(fldr, filename + "_" + str(idx) + ".csv")
while os.path.exists(path):
idx += 1
path = os.path.join(fldr, filename + "_" + str(idx) + ".csv")
pd.DataFrame(data).to_csv(path, header=None, index=None)
@property
def _torso_xpos(self):
"""
Grabs torso center position
Returns:
np.array: torso pos (x,y,z)
"""
return np.array(self.sim.data.body_xpos[self.torso_body_id])
class Reach(SingleArmEnv):
"""
This class corresponds to the reaching task for a single robot arm.
Args:
robots (str or list of str): Specification for specific robot arm(s) to be instantiated within this env
(e.g: "Sawyer" would generate one arm; ["Panda", "Panda", "Sawyer"] would generate three robot arms)
Note: Must be a single single-arm robot!
env_configuration (str): Specifies how to position the robots within the environment (default is "default").
For most single arm environments, this argument has no impact on the robot setup.
controller_configs (str or list of dict): If set, contains relevant controller parameters for creating a
custom controller. Else, uses the default controller for this specific task. Should either be single
dict if same controller is to be used for all robots or else it should be a list of the same length as
"robots" param
gripper_types (str or list of str): type of gripper, used to instantiate
gripper models from gripper factory. Default is "default", which is the default grippers(s) associated
with the robot(s) the 'robots' specification. None removes the gripper, and any other (valid) model
overrides the default gripper. Should either be single str if same gripper type is to be used for all
robots or else it should be a list of the same length as "robots" param
initialization_noise (dict or list of dict): Dict containing the initialization noise parameters.
The expected keys and corresponding value types are specified below:
:`'magnitude'`: The scale factor of uni-variate random noise applied to each of a robot's given initial
joint positions. Setting this value to `None` or 0.0 results in no noise being applied.
If "gaussian" type of noise is applied then this magnitude scales the standard deviation applied,
If "uniform" type of noise is applied then this magnitude sets the bounds of the sampling range
:`'type'`: Type of noise to apply. Can either specify "gaussian" or "uniform"
Should either be single dict if same noise value is to be used for all robots or else it should be a
list of the same length as "robots" param
:Note: Specifying "default" will automatically use the default noise settings.
Specifying None will automatically create the required dict with "magnitude" set to 0.0.
table_full_size (3-tuple): x, y, and z dimensions of the table.
table_friction (3-tuple): the three mujoco friction parameters for
the table.
use_camera_obs (bool): if True, every observation includes rendered image(s)
use_object_obs (bool): if True, include object (ball) information in
the observation.
reward_scale (None or float): Scales the normalized reward function by the amount specified.
If None, environment reward remains unnormalized
reward_shaping (bool): if True, use dense rewards.
placement_initializer (ObjectPositionSampler): if provided, will
be used to place objects on every reset, else a UniformRandomSampler
is used by default.
has_renderer (bool): If true, render the simulation state in
a viewer instead of headless mode.
has_offscreen_renderer (bool): True if using off-screen rendering
render_camera (str): Name of camera to render if `has_renderer` is True. Setting this value to 'None'
will result in the default angle being applied, which is useful as it can be dragged / panned by
the user using the mouse
render_collision_mesh (bool): True if rendering collision meshes in camera. False otherwise.
render_visual_mesh (bool): True if rendering visual meshes in camera. False otherwise.
render_gpu_device_id (int): corresponds to the GPU device id to use for offscreen rendering.
Defaults to -1, in which case the device will be inferred from environment variables
(GPUS or CUDA_VISIBLE_DEVICES).
control_freq (float): how many control signals to receive in every second. This sets the amount of
simulation time that passes between every action input.
horizon (int): Every episode lasts for exactly @horizon timesteps.
ignore_done (bool): True if never terminating the environment (ignore @horizon).
hard_reset (bool): If True, re-loads model, sim, and render object upon a reset call, else,
only calls sim.reset and resets all robosuite-internal variables
camera_names (str or list of str): name of camera to be rendered. Should either be single str if
same name is to be used for all cameras' rendering or else it should be a list of cameras to render.
:Note: At least one camera must be specified if @use_camera_obs is True.
:Note: To render all robots' cameras of a certain type (e.g.: "robotview" or "eye_in_hand"), use the
convention "all-{name}" (e.g.: "all-robotview") to automatically render all camera images from each
robot's camera list).
camera_heights (int or list of int): height of camera frame. Should either be single int if
same height is to be used for all cameras' frames or else it should be a list of the same length as
"camera names" param.
camera_widths (int or list of int): width of camera frame. Should either be single int if
same width is to be used for all cameras' frames or else it should be a list of the same length as
"camera names" param.
camera_depths (bool or list of bool): True if rendering RGB-D, and RGB otherwise. Should either be single
bool if same depth setting is to be used for all cameras or else it should be a list of the same length as
"camera names" param.
camera_segmentations (None or str or list of str or list of list of str): Camera
for each camera. Valid options are:
`None`: no segmentation sensor used
`'instance'`: segmentation at the class-instance level
`'class'`: segmentation at the class level
`'element'`: segmentation at the per-geom level
If not None, multiple types of segmentations can be specified. A [list of st
[multiple / a single] segmentation(s) to use for all cameras. A list of list
segmentation setting(s) to use.
max_object_distance: max distance in meters to randomize object from end effector
Raises:
AssertionError: [Invalid number of robots specified]
"""
def __init__(
self,
robots,
env_configuration="default",
controller_configs=None,
gripper_types="default",
initialization_noise="default",
table_full_size=(0.8, 0.8, 0.05),
table_friction=(1., 5e-3, 1e-4),
use_camera_obs=True,
use_object_obs=True,
reward_scale=1.0,
reward_shaping=False,
placement_initializer=None,
has_renderer=False,
has_offscreen_renderer=True,
render_camera="frontview",
render_collision_mesh=False,
render_visual_mesh=True,
render_gpu_device_id=-1,
control_freq=20,
horizon=1000,
ignore_done=False,
hard_reset=True,
camera_names="agentview",
camera_heights=256,
camera_widths=256,
camera_depths=False,
camera_segmentations=None,
renderer="mujoco",
renderer_config=None,
max_object_distance=0.2,
):
self.max_object_distance = np.abs(max_object_distance)
# settings for table top
self.table_full_size = table_full_size
self.table_friction = table_friction
self.table_offset = np.array((0, 0, 0.8))
# reward configuration
self.reward_scale = reward_scale
self.reward_shaping = reward_shaping
# whether to use ground-truth object states
self.use_object_obs = use_object_obs
# object placement initializer
self.placement_initializer = placement_initializer
super().__init__(
robots=robots,
env_configuration=env_configuration,
controller_configs=controller_configs,
mount_types="default",
gripper_types=gripper_types,
initialization_noise=initialization_noise,
use_camera_obs=use_camera_obs,
has_renderer=has_renderer,
has_offscreen_renderer=has_offscreen_renderer,
render_camera=render_camera,
render_collision_mesh=render_collision_mesh,
render_visual_mesh=render_visual_mesh,
render_gpu_device_id=render_gpu_device_id,
control_freq=control_freq,
horizon=horizon,
ignore_done=ignore_done,
hard_reset=hard_reset,
camera_names=camera_names,
camera_heights=camera_heights,
camera_widths=camera_widths,
camera_depths=camera_depths,
camera_segmentations=camera_segmentations,
renderer=renderer,
renderer_config=renderer_config,
)
def reward(self, action=None):
"""
Reward function for the task.
Sparse un-normalized reward:
- a discrete reward of 2.25 is provided if the ball is grasped
Un-normalized summed components if using reward shaping:
- Reaching: in [0, 1], to encourage the arm to reach the ball
- Grasping: in {0, 1}, non-zero if arm is grasping the ball
The sparse reward only consists of the reach component.
Note that the final reward is normalized and scaled by
reward_scale / 2.25 as well so that the max score is equal to reward_scale
Args:
action (np array): [NOT USED]
Returns:
float: reward value
"""
reward = 0.
# sparse completion reward
if self._check_success():
reward = 2.25
# use a shaping reward
elif self.reward_shaping:
# reaching reward
ball_pos = self.sim.data.body_xpos[self.ball_body_id]
gripper_site_pos = self.sim.data.site_xpos[
self.robots[0].eef_site_id]
dist = np.linalg.norm(gripper_site_pos - ball_pos)
reaching_reward = 1 - np.tanh(10.0 * dist)
reward += reaching_reward
# Scale reward if requested
if self.reward_scale is not None:
reward *= self.reward_scale / 2.25
return reward
def _load_model(self):
"""
Loads an xml model, puts it in self.model
"""
super()._load_model()
# Adjust base pose accordingly
xpos = self.robots[0].robot_model.base_xpos_offset["table"](
self.table_full_size[0])
self.robots[0].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,
#)
mujoco_arena = EmptyArena()
# Arena always gets set to zero origin
mujoco_arena.set_origin([0, 0, 0])
# initialize objects of interest
tex_attrib = {
"type": "ball",
}
mat_attrib = {
"texrepeat": "1 1",
"specular": "0.4",
"shininess": "0.1",
}
redwood = CustomMaterial(
texture="WoodRed",
tex_name="redwood",
mat_name="redwood_mat",
tex_attrib=tex_attrib,
mat_attrib=mat_attrib,
)
self.ball = BallObject(
name="ball",
size_min=[0.018], # [0.015, 0.015, 0.015],
size_max=[0.024], # [0.018, 0.018, 0.018])
rgba=[1, 0, 0, 1],
density=1,
joints=[
{"name":"ball_x",
"type":"slide",
"damping":100,
"axis":"1 0 0",
},
{"name":"ball_y",
"type":"slide",
"damping":100,
"axis":"0 1 0"},
{"name":"ball_z",
"type":"slide",
"axis":"0 0 1",
"damping":100,
}
]
)
# get current end effector position
#max_pos = [-.75, .75] #np.clip(self.max_object_distance + eef_pos[:3], -1, 1)
#min_pos = [-.75, .75] #np.clip(-self.max_object_distance + eef_pos[:3], -1, 1)
## Create placement initializer
# TODO I think we need to create every time since joint position will be different
if self.placement_initializer is not None:
self.placement_initializer.reset()
self.placement_initializer.add_objects(self.ball)
else:
self.placement_initializer = UniformRandomSampler(
name="ObjectSampler",
mujoco_objects=self.ball,
x_range=[-.75, .75],
y_range=[-.75, .75],
rotation=None,
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.ball,
)
def _setup_references(self):
"""
Sets up references to important components. A reference is typically an
index or a list of indices that point to the corresponding elements
in a flatten array, which is how MuJoCo stores physical simulation data.
"""
super()._setup_references()
# Additional object references from this env
self.ball_body_id = self.sim.model.body_name2id(self.ball.root_body)
def _setup_observables(self):
"""
Sets up observables to be used for this environment. Creates object-based observables if enabled
Returns:
OrderedDict: Dictionary mapping observable names to its corresponding Observable object
"""
observables = super()._setup_observables()
# low-level object information
if self.use_object_obs:
# Get robot prefix and define observables modality
pf = self.robots[0].robot_model.naming_prefix
modality = "object"
# ball-related observables
@sensor(modality=modality)
def ball_pos(obs_cache):
return np.array(self.sim.data.body_xpos[self.ball_body_id])
@sensor(modality=modality)
def ball_quat(obs_cache):
return convert_quat(np.array(
self.sim.data.body_xquat[self.ball_body_id]),
to="xyzw")
@sensor(modality=modality)
def gripper_to_ball_pos(obs_cache):
return obs_cache[f"{pf}eef_pos"] - obs_cache["ball_pos"] if \
f"{pf}eef_pos" in obs_cache and "ball_pos" in obs_cache else np.zeros(3)
sensors = [ball_pos, ball_quat, gripper_to_ball_pos]
names = [s.__name__ for s in sensors]
# Create observables
for name, s in zip(names, sensors):
observables[name] = Observable(
name=name,
sensor=s,
sampling_rate=self.control_freq,
)
return observables
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:
eef_pos = self.robots[0].sim.data.get_body_xpos('gripper0_eef')
self.ball_initializer = UniformRandomSampler(
name="ObjectSampler",
mujoco_objects=self.ball,
x_range=[-self.max_object_distance, self.max_object_distance],
y_range=[-self.max_object_distance, self.max_object_distance],
rotation=None,
reference_pos=eef_pos,
ensure_object_boundary_in_range=False,
ensure_valid_placement=True,
)
ball_id = self.sim.model.body_name2id(self.ball.root_body)
distance = 1.0
print(distance)
while distance > .8:
# TODO what about z?
z = np.random.uniform(high=eef_pos[2]+self.max_object_distance, low=eef_pos[2]-self.max_object_distance)
# Sample from the placement initializer for all objects
object_placements = self.ball_initializer.sample()
new_position = np.array(object_placements['ball'][0])
new_position[2] = z
distance = np.sqrt(np.sum(new_position)**2)
print('setting ball', new_position)
self.sim.data.body_xpos[ball_id] = new_position
# below only works if there are joints - but we don't want joint bc we need static ball
# Loop through all objects and reset their positions
self.sim.data.set_joint_qpos('ball_x', new_position[0])
self.sim.data.set_joint_qpos('ball_y', new_position[1])
self.sim.data.set_joint_qpos('ball_z', new_position[2])
#for obj_pos, obj_quat, obj in object_placements.values():
# self.sim.data.set_joint_qpos(
# obj.joints[0],
# np.concatenate([np.array(obj_pos),
# np.array(obj_quat)]))
def visualize(self, vis_settings):
"""
In addition to super call, visualize gripper site proportional to the distance to the ball.
Args:
vis_settings (dict): Visualization keywords mapped to T/F, determining whether that specific
component should be visualized. Should have "grippers" keyword as well as any other relevant
options specified.
"""
# Run superclass method first
super().visualize(vis_settings=vis_settings)
# Color the gripper visualization site according to its distance to the ball
if vis_settings["grippers"]:
self._visualize_gripper_to_target(gripper=self.robots[0].gripper,
target=self.ball)
def _check_success(self):
"""
Check if ball has been grasped
Returns:
bool: True if ball has been grasped
"""
# grasping reward
if self._check_grasp(gripper=self.robots[0].gripper,
object_geoms=self.ball):
return True
else:
return False