class GripperTester:
"""
A class that is used to test gripper
Args:
gripper (GripperModel): A gripper instance to be tested
pos (str): (x y z) position to place the gripper in string form, e.g. '0 0 0.3'
quat (str): rotation to apply to gripper in string form, e.g. '0 0 1 0' to flip z axis
gripper_low_pos (float): controls the gipper y position, larger -> higher
gripper_high_pos (float): controls the gipper y high position larger -> higher,
must be larger than gripper_low_pos
box_size (None or 3-tuple of int): the size of the box to grasp, None defaults to [0.02, 0.02, 0.02]
box_density (int): the density of the box to grasp
step_time (int): the interval between two gripper actions
render (bool): if True, show rendering
"""
def __init__(
self,
gripper,
pos,
quat,
gripper_low_pos,
gripper_high_pos,
box_size=None,
box_density=10000,
step_time=400,
render=True,
):
# define viewer
self.viewer = None
world = MujocoWorldBase()
# Add a table
arena = TableArena(table_full_size=(0.4, 0.4, 0.1), table_offset=(0, 0, 0.1), has_legs=False)
world.merge(arena)
# Add a gripper
self.gripper = gripper
# Create another body with a slider joint to which we'll add this gripper
gripper_body = ET.Element("body")
gripper_body.set("pos", pos)
gripper_body.set("quat", quat) # flip z
gripper_body.append(new_joint(name="gripper_z_joint", type="slide", axis="0 0 -1", damping="50"))
# Add all gripper bodies to this higher level body
for body in gripper.worldbody:
gripper_body.append(body)
# Merge the all of the gripper tags except its bodies
world.merge(gripper, merge_body=None)
# Manually add the higher level body we created
world.worldbody.append(gripper_body)
# Create a new actuator to control our slider joint
world.actuator.append(new_actuator(joint="gripper_z_joint", act_type="position", name="gripper_z", kp="500"))
# Add an object for grasping
# density is in units kg / m3
TABLE_TOP = [0, 0, 0.09]
if box_size is None:
box_size = [0.02, 0.02, 0.02]
box_size = np.array(box_size)
self.cube = BoxObject(
name="object", size=box_size, rgba=[1, 0, 0, 1], friction=[1, 0.005, 0.0001], density=box_density
)
object_pos = np.array(TABLE_TOP + box_size * [0, 0, 1])
mujoco_object = self.cube.get_obj()
# Set the position of this object
mujoco_object.set("pos", array_to_string(object_pos))
# Add our object to the world body
world.worldbody.append(mujoco_object)
# add reference objects for x and y axes
x_ref = BoxObject(
name="x_ref", size=[0.01, 0.01, 0.01], rgba=[0, 1, 0, 1], obj_type="visual", joints=None
).get_obj()
x_ref.set("pos", "0.2 0 0.105")
world.worldbody.append(x_ref)
y_ref = BoxObject(
name="y_ref", size=[0.01, 0.01, 0.01], rgba=[0, 0, 1, 1], obj_type="visual", joints=None
).get_obj()
y_ref.set("pos", "0 0.2 0.105")
world.worldbody.append(y_ref)
self.world = world
self.render = render
self.simulation_ready = False
self.step_time = step_time
self.cur_step = 0
if gripper_low_pos > gripper_high_pos:
raise ValueError(
"gripper_low_pos {} is larger " "than gripper_high_pos {}".format(gripper_low_pos, gripper_high_pos)
)
self.gripper_low_pos = gripper_low_pos
self.gripper_high_pos = gripper_high_pos
def start_simulation(self):
"""
Starts simulation of the test world
"""
model = self.world.get_model(mode="mujoco_py")
self.sim = MjSim(model)
if self.render:
self.viewer = MjViewer(self.sim)
self.sim_state = self.sim.get_state()
# For gravity correction
gravity_corrected = ["gripper_z_joint"]
self._gravity_corrected_qvels = [self.sim.model.get_joint_qvel_addr(x) for x in gravity_corrected]
self.gripper_z_id = self.sim.model.actuator_name2id("gripper_z")
self.gripper_z_is_low = False
self.gripper_actuator_ids = [self.sim.model.actuator_name2id(x) for x in self.gripper.actuators]
self.gripper_is_closed = True
self.object_id = self.sim.model.body_name2id(self.cube.root_body)
object_default_pos = self.sim.data.body_xpos[self.object_id]
self.object_default_pos = np.array(object_default_pos, copy=True)
self.reset()
self.simulation_ready = True
def reset(self):
"""
Resets the simulation to the initial state
"""
self.sim.set_state(self.sim_state)
self.cur_step = 0
def close(self):
"""
Close the viewer if it exists
"""
if self.viewer is not None:
self.viewer.close()
def step(self):
"""
Forward the simulation by one timestep
Raises:
RuntimeError: if start_simulation is not yet called.
"""
if not self.simulation_ready:
raise RuntimeError("Call start_simulation before calling step")
if self.gripper_z_is_low:
self.sim.data.ctrl[self.gripper_z_id] = self.gripper_low_pos
else:
self.sim.data.ctrl[self.gripper_z_id] = self.gripper_high_pos
if self.gripper_is_closed:
self._apply_gripper_action(1)
else:
self._apply_gripper_action(-1)
self._apply_gravity_compensation()
self.sim.step()
if self.render:
self.viewer.render()
self.cur_step += 1
def _apply_gripper_action(self, action):
"""
Applies binary gripper action
Args:
action (int): Action to apply. Should be -1 (open) or 1 (closed)
"""
gripper_action_actual = self.gripper.format_action(np.array([action]))
# rescale normalized gripper action to control ranges
ctrl_range = self.sim.model.actuator_ctrlrange[self.gripper_actuator_ids]
bias = 0.5 * (ctrl_range[:, 1] + ctrl_range[:, 0])
weight = 0.5 * (ctrl_range[:, 1] - ctrl_range[:, 0])
applied_gripper_action = bias + weight * gripper_action_actual
self.sim.data.ctrl[self.gripper_actuator_ids] = applied_gripper_action
def _apply_gravity_compensation(self):
"""
Applies gravity compensation to the simulation
"""
self.sim.data.qfrc_applied[self._gravity_corrected_qvels] = self.sim.data.qfrc_bias[
self._gravity_corrected_qvels
]
def loop(self, total_iters=1, test_y=False, y_baseline=0.01):
"""
Performs lower, grip, raise and release actions of a gripper,
each separated with T timesteps
Args:
total_iters (int): Iterations to perform before exiting
test_y (bool): test if object is lifted
y_baseline (float): threshold for determining that object is lifted
"""
seq = [(False, False), (True, False), (True, True), (False, True)]
for cur_iter in range(total_iters):
for cur_plan in seq:
self.gripper_z_is_low, self.gripper_is_closed = cur_plan
for step in range(self.step_time):
self.step()
if test_y:
if not self.object_height > y_baseline:
raise ValueError(
"object is lifed by {}, ".format(self.object_height)
+ "not reaching the requirement {}".format(y_baseline)
)
@property
def object_height(self):
"""
Queries the height (z) of the object compared to on the ground
Returns:
float: Object height relative to default (ground) object position
"""
return self.sim.data.body_xpos[self.object_id][2] - self.object_default_pos[2]
class MujocoEnv(gym.Env):
"""Initializes a Mujoco Environment."""
def __init__(self,
has_renderer=False,
has_offscreen_renderer=True,
render_collision_mesh=False,
render_visual_mesh=True,
control_freq=10,
horizon=1000,
ignore_done=False,
use_camera_obs=False,
camera_name="frontview",
camera_height=256,
camera_width=256,
camera_depth=False,
**kwargs):
"""
Args:
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_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.
control_freq (float): how many control signals to receive
in every simulated 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).
use_camera_obs (bool): if True, every observation includes a
rendered image.
camera_name (str): name of camera to be rendered. Must be
set if @use_camera_obs is True.
camera_height (int): height of camera frame.
camera_width (int): width of camera frame.
camera_depth (bool): True if rendering RGB-D, and RGB otherwise.
"""
self.seed()
self._kwargs = kwargs
self._screen_width = kwargs["screen_width"]
self._screen_height = kwargs["screen_height"]
self.has_renderer = has_renderer
self.has_offscreen_renderer = has_offscreen_renderer
self.render_collision_mesh = render_collision_mesh
self.render_visual_mesh = render_visual_mesh
self.control_freq = control_freq
self.horizon = horizon
self.ignore_done = ignore_done
self.viewer = None
self.model = None
# settings for camera observations
self.use_camera_obs = use_camera_obs
if self.use_camera_obs and not self.has_offscreen_renderer:
raise ValueError(
"Camera observations require an offscreen renderer.")
self.camera_name = camera_name
if self.use_camera_obs and self.camera_name is None:
raise ValueError("Must specify camera name when using camera obs")
self.camera_height = camera_height
self.camera_width = camera_width
self.camera_depth = camera_depth
self.pid = None
self._reset_internal()
def initialize_time(self, control_freq):
"""
Initializes the time constants used for simulation.
"""
self.cur_time = 0
self.model_timestep = self.sim.model.opt.timestep
if self.model_timestep <= 0:
raise XMLError("xml model defined non-positive time step")
self.control_freq = control_freq
if control_freq <= 0:
raise SimulationError(
"control frequency {} is invalid".format(control_freq))
self.control_timestep = 1. / control_freq
def _setup_pid(self):
"Function to setup the pid controller. Should set the self.pid attribute."
raise NotImplementedError
def _load_model(self):
"""Loads an xml model, puts it in self.model"""
pass
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.
"""
pass
def reset(self):
"""Resets simulation."""
# if there is an active viewer window, destroy it
self._destroy_viewer()
self._reset_internal()
self.sim.forward()
return self._get_observation()
def _reset_internal(self):
"""Resets simulation internal configurations."""
# instantiate simulation from MJCF model
self._load_model()
self.mjpy_model = self.model.get_model(mode="mujoco_py")
self.sim = MjSim(self.mjpy_model)
self.initialize_time(self.control_freq)
# create visualization screen or renderer
if self.has_renderer and self.viewer is None:
self.viewer = MjViewer(self.sim)
self.viewer.viewer.vopt.geomgroup[0] = (
1 if self.render_collision_mesh else 0)
self.viewer.viewer.vopt.geomgroup[
1] = 1 if self.render_visual_mesh else 0
# hiding the overlay speeds up rendering significantly
self.viewer.viewer._hide_overlay = True
self.viewer.viewer._render_every_frame = True
elif self.has_offscreen_renderer:
if self.sim._render_context_offscreen is None:
render_context = MjRenderContextOffscreen(self.sim)
self.sim.add_render_context(render_context)
self.sim._render_context_offscreen.vopt.geomgroup[0] = (
1 if self.render_collision_mesh else 0)
self.sim._render_context_offscreen.vopt.geomgroup[1] = (
1 if self.render_visual_mesh else 0)
# additional housekeeping
self.sim_state_initial = self.sim.get_state()
self._get_reference()
self.cur_time = 0
self.timestep = 0
self.done = False
def _get_observation(self):
"""Returns an OrderedDict containing observations [(name_string, np.array), ...]."""
return OrderedDict()
def step(self, action):
"""Takes a step in simulation with control command @action."""
if self.done:
raise ValueError("executing action in terminated episode")
self.timestep += 1
self._pre_action(action)
end_time = self.cur_time + self.control_timestep
while self.cur_time < end_time:
if (self.pid is not None):
self._set_pid_control()
self.sim.step()
self.cur_time += self.model_timestep
reward, done, info = self._post_action(action)
return self._get_observation(), reward, done, info
def _set_pid_control(self):
"Do any processing required with the pid"
current_qvel = self.sim.data.qvel
self.sim.data.ctrl[:] = self.pid(current_qvel, self.model_timestep)
def _pre_action(self, action):
"""Do any preprocessing before taking an action."""
self.sim.data.ctrl[:] = action
def _post_action(self, action):
"""Do any housekeeping after taking an action."""
reward = self.reward(action)
# done if number of elapsed timesteps is greater than horizon
self.done = (self.timestep >= self.horizon) and not self.ignore_done
return reward, self.done, {}
def reward(self, action):
"""Reward should be a function of state and action."""
return 0
def render(self, mode="rgb_array", height=100, width=100):
"""
Renders to an off-screen window.
"""
if height == None:
height = self._screen_height
if width == None:
width = self._screen_width
if mode == "rgb_array":
camera_obs = self.sim.render(
camera_name="birdview",
width=width,
height=height,
depth=False,
)
#camera_obs = camera_obs[::-1, :, :] / 255.0
assert np.sum(camera_obs) > 0, "rendering image is blank"
return camera_obs
elif mode == "human":
self.viewer().render()
return None
def seed(self, seed=None):
self.np_random, seed = seeding.np_random(seed)
return [seed]
def observation_spec(self):
"""
Returns an observation as observation specification.
An alternative design is to return an OrderedDict where the keys
are the observation names and the values are the shapes of observations.
We leave this alternative implementation commented out, as we find the
current design is easier to use in practice.
"""
observation = self._get_observation()
# return observation
observation_spec = OrderedDict()
for k, v in observation.items():
observation_spec[k] = v.shape
return observation_spec
def action_spec(self):
"""
Action specification should be implemented in subclasses.
Action space is represented by a tuple of (low, high), which are two numpy
vectors that specify the min/max action limits per dimension.
"""
raise NotImplementedError
def reset_from_xml_string(self, xml_string):
"""Reloads the environment from an XML description of the environment."""
# if there is an active viewer window, destroy it
self.close()
# load model from xml
self.mjpy_model = load_model_from_xml(xml_string)
self.sim = MjSim(self.mjpy_model)
self.initialize_time(self.control_freq)
if self.has_renderer and self.viewer is None:
self.viewer = MjViewer(self.sim)
self.viewer.viewer.vopt.geomgroup[0] = (
1 if self.render_collision_mesh else 0)
self.viewer.viewer.vopt.geomgroup[
1] = 1 if self.render_visual_mesh else 0
# hiding the overlay speeds up rendering significantly
self.viewer.viewer._hide_overlay = True
elif self.has_offscreen_renderer:
render_context = MjRenderContextOffscreen(self.sim)
render_context.vopt.geomgroup[
0] = 1 if self.render_collision_mesh else 0
render_context.vopt.geomgroup[
1] = 1 if self.render_visual_mesh else 0
self.sim.add_render_context(render_context)
self.sim_state_initial = self.sim.get_state()
self._get_reference()
self.cur_time = 0
self.timestep = 0
self.done = False
# necessary to refresh MjData
self.sim.forward()
def find_contacts(self, geoms_1, geoms_2):
"""
Finds contact between two geom groups.
Args:
geoms_1: a list of geom names (string)
geoms_2: another list of geom names (string)
Returns:
iterator of all contacts between @geoms_1 and @geoms_2
"""
for contact in self.sim.data.contact[0:self.sim.data.ncon]:
# check contact geom in geoms
c1_in_g1 = self.sim.model.geom_id2name(contact.geom1) in geoms_1
c2_in_g2 = self.sim.model.geom_id2name(contact.geom2) in geoms_2
# check contact geom in geoms (flipped)
c2_in_g1 = self.sim.model.geom_id2name(contact.geom2) in geoms_1
c1_in_g2 = self.sim.model.geom_id2name(contact.geom1) in geoms_2
if (c1_in_g1 and c2_in_g2) or (c1_in_g2 and c2_in_g1):
yield contact
def _check_contact(self):
"""Returns True if gripper is in contact with an object."""
return False
def _check_success(self):
"""
Returns True if task has been completed.
"""
return False
def _destroy_viewer(self):
# if there is an active viewer window, destroy it
if self.viewer is not None:
self.viewer.close() # change this to viewer.finish()?
self.viewer = None
def close(self):
"""Do any cleanup necessary here."""
self._destroy_viewer()
def get_state(self):
return self.sim.get_state()
def set_state(self, qpos, qvel):
assert qpos.shape == (self.sim.model.nq, ) and qvel.shape == (
self.sim.model.nv, )
old_state = self.sim.get_state()
new_state = mujoco_py.MjSimState(old_state.time, qpos, qvel,
old_state.act, old_state.udd_state)
self.sim.set_state(new_state)
self.sim.forward()
def renderer_on(self):
self.has_renderer = True
def renderer_off(self):
self.has_renderer = False
