def _initialize_sim(self, xml_string=None):
"""
Creates a MjSim object and stores it in self.sim. If @xml_string is specified, the MjSim object will be created
from the specified xml_string. Else, it will pull from self.model to instantiate the simulation
Args:
xml_string (str): If specified, creates MjSim object from this filepath
"""
super()._initialize_sim(xml_string)
if self._use_dm_backend:
if not hasattr(self, "camera_settings"):
self.camera_settings = {}
with io.StringIO() as string:
string.write(ET.tostring(self.model.root, encoding="unicode"))
st = string.getvalue()
dm_mujoco = module.get_dm_mujoco()
self.dm_sim = dm_mujoco.Physics.from_xml_string(st)
patch_mjlib_accessors(
module.get_dm_mujoco().wrapper.mjbindings.mjlib,
self.dm_sim.model,
self.dm_sim.data,
)
self.renderer = DMRenderer(
self.dm_sim,
clear_geom_group_0=True,
camera_settings=self.camera_settings,
camera_select_next=False,
mjpy_sim=self.sim,
)
def render_offscreen(self,
width: int,
height: int,
mode: RenderMode = RenderMode.RGB,
camera_id: int = -1) -> np.ndarray:
"""Renders the camera view as a NumPy array of pixels.
Args:
width: The viewport width (pixels).
height: The viewport height (pixels).
mode: The rendering mode.
camera_id: The ID of the camera to render from. By default, uses
the free camera.
Returns:
A NumPy array of the pixels.
"""
mujoco = module.get_dm_mujoco()
# TODO(michaelahn): Consider caching the camera.
camera = mujoco.Camera(physics=self._physics,
height=height,
width=width,
camera_id=camera_id)
# Update the camera configuration for the free-camera.
if camera_id == -1:
self._update_camera(
camera._render_camera, # pylint: disable=protected-access
)
image = camera.render(depth=(mode == RenderMode.DEPTH),
segmentation=(mode == RenderMode.SEGMENTATION))
camera._scene.free() # pylint: disable=protected-access
return image
def update(self, force=False):
"""
Updates the state of the robot arm, including end effector pose / orientation / velocity, joint pos/vel,
jacobian, and mass matrix. By default, since this is a non-negligible computation, multiple redundant calls
will be ignored via the self.new_update attribute flag. However, if the @force flag is set, the update will
occur regardless of that state of self.new_update. This base class method of @run_controller resets the
self.new_update flag
Args:
force (bool): Whether to force an update to occur or not
"""
# Only run update if self.new_update or force flag is set
if self.new_update or force:
self.sim.forward()
self.ee_pos = np.array(
self.sim.data.site_xpos[self.sim.model.site_name2id(
self.eef_name)])
self.ee_ori_mat = np.array(
self.sim.data.site_xmat[self.sim.model.site_name2id(
self.eef_name)].reshape([3, 3]))
self.ee_pos_vel = np.array(
self.sim.data.site_xvelp[self.sim.model.site_name2id(
self.eef_name)])
self.ee_ori_vel = np.array(
self.sim.data.site_xvelr[self.sim.model.site_name2id(
self.eef_name)])
self.joint_pos = np.array(self.sim.data.qpos[self.qpos_index])
self.joint_vel = np.array(self.sim.data.qvel[self.qvel_index])
self.J_pos = np.array(
self.sim.data.get_site_jacp(self.eef_name).reshape(
(3, -1))[:, self.qvel_index])
self.J_ori = np.array(
self.sim.data.get_site_jacr(self.eef_name).reshape(
(3, -1))[:, self.qvel_index])
self.J_full = np.array(np.vstack([self.J_pos, self.J_ori]))
mass_matrix = np.ndarray(shape=(len(self.sim.data.qvel)**2, ),
dtype=np.float64,
order='C')
#TODO: definitely clean this up and subclass
try:
mujoco_py.cymj._mj_fullM(self.sim.model, mass_matrix,
self.sim.data.qM)
except:
mjlib = module.get_dm_mujoco().wrapper.mjbindings.mjlib
mjlib.mj_fullM(self.sim.model.ptr, mass_matrix,
self.sim.data.qM)
mass_matrix = np.reshape(
mass_matrix,
(len(self.sim.data.qvel), len(self.sim.data.qvel)))
self.mass_matrix = mass_matrix[self.qvel_index, :][:,
self.qvel_index]
# Clear self.new_update
self.new_update = False
def __init__(self, physics, **kwargs):
assert isinstance(physics, module.get_dm_mujoco().Physics), \
'DMRenderer takes a DM Control Physics object.'
super().__init__(**kwargs)
self._physics = physics
self._window = None
# Set the camera to lookat the center of the geoms. (mujoco_py does
# this automatically.
if 'lookat' not in self._camera_settings:
self._camera_settings['lookat'] = [
np.median(self._physics.data.geom_xpos[:, i]) for i in range(3)
]
def __init__(
self,
model_file: str,
use_dm_backend: bool = False,
camera_settings: Optional[Dict] = None,
):
"""Initializes a new simulation.
Args:
model_file: The MuJoCo XML model file to load.
use_dm_backend: If True, uses DM Control's Physics (MuJoCo v2.0) as
the backend for the simulation. Otherwise, uses mujoco_py (MuJoCo
v1.5) as the backend.
camera_settings: Settings to initialize the renderer's camera. This
can contain the keys `distance`, `azimuth`, and `elevation`.
"""
self._use_dm_backend = use_dm_backend
if not os.path.isfile(model_file):
raise ValueError(
"[MujocoSimRobot] Invalid model file path: {}".format(
model_file))
if self._use_dm_backend:
dm_mujoco = module.get_dm_mujoco()
if model_file.endswith(".mjb"):
self.sim = dm_mujoco.Physics.from_binary_path(model_file)
else:
self.sim = dm_mujoco.Physics.from_xml_path(model_file)
self.model = self.sim.model
self._patch_mjlib_accessors(self.model, self.sim.data)
self.renderer = DMRenderer(self.sim,
camera_settings=camera_settings)
else: # Use mujoco_py
mujoco_py = module.get_mujoco_py()
self.model = mujoco_py.load_model_from_path(model_file)
self.sim = mujoco_py.MjSim(self.model)
self.renderer = MjPyRenderer(self.sim,
camera_settings=camera_settings)
self.data = self.sim.data
def get_mjlib(self):
"""Returns an object that exposes the low-level MuJoCo API."""
if self._use_dm_backend:
return module.get_dm_mujoco().wrapper.mjbindings.mjlib
else:
return module.get_mujoco_py_mjlib()
def get_mjlib(self):
"""Returns an object that exposes the low-level MuJoCo API."""
return module.get_dm_mujoco().wrapper.mjbindings.mjlib
def __init__(self, model_path, frame_skip, rgb_array_res=(640, 480)):
if not path.exists(model_path):
raise IOError("File %s does not exist" % model_path)
if self.control_mode == "vices":
if model_path == (
"/home/mdalal/research/metaworld/metaworld/envs/assets_v2/sawyer_xyz/sawyer_basketball.xml"
):
model_path = "/home/mdalal/research/metaworld/metaworld/envs/assets_v2/sawyer_xyz/sawyer_basketball_torque.xml"
self.reset_qpos = np.array([
0.00000000e00,
6.00000000e-01,
2.98721632e-02,
1.00000000e00,
0.00000000e00,
0.00000000e00,
0.00000000e00,
1.88500731e00,
-5.88898933e-01,
-2.50919766e-02,
6.95071420e-01,
2.99999993e-02,
1.02969815e00,
2.31183042e00,
-1.71909704e-04,
1.71743732e-04,
])
elif model_path == (
"/home/mdalal/research/metaworld/metaworld/envs/assets_v2/sawyer_xyz/sawyer_drawer.xml"
):
model_path = "/home/mdalal/research/metaworld/metaworld/envs/assets_v2/sawyer_xyz/sawyer_drawer_torque.xml"
self.reset_qpos = [
1.88500731e00,
-5.88898933e-01,
-9.64183847e-01,
1.64749509e00,
9.28632075e-01,
1.02969815e00,
2.31183042e00,
-1.71909704e-04,
1.71743732e-04,
-1.50000000e-01,
]
elif model_path == (
"/home/mdalal/research/metaworld/metaworld/envs/assets_v2/sawyer_xyz/sawyer_soccer.xml"
):
model_path = "/home/mdalal/research/metaworld/metaworld/envs/assets_v2/sawyer_xyz/sawyer_soccer_torque.xml"
self.reset_qpos = [
1.88500734e00,
-5.88898932e-01,
-9.64183883e-01,
1.64749516e00,
9.28632122e-01,
1.02969813e00,
2.31183036e00,
-1.71909704e-04,
1.71743732e-04,
-8.83832789e-02,
6.86617607e-01,
3.00000000e-02,
1.00000000e00,
0.00000000e00,
0.00000000e00,
0.00000000e00,
]
elif model_path == (
"/home/mdalal/research/metaworld/metaworld/envs/assets_v2/sawyer_xyz/sawyer_table_with_hole.xml"
):
model_path = "/home/mdalal/research/metaworld/metaworld/envs/assets_v2/sawyer_xyz/sawyer_table_with_hole_torque.xml"
self.reset_qpos = [
1.88500734e00,
-5.88898932e-01,
-9.64183883e-01,
1.64749516e00,
9.28632122e-01,
1.02969813e00,
2.31183036e00,
-1.71909704e-04,
1.71743732e-04,
-8.83832789e-02,
6.86617607e-01,
6.99354863e-02,
1.00000000e00,
0.00000000e00,
0.00000000e00,
0.00000000e00,
]
elif model_path == (
"/home/mdalal/research/metaworld/metaworld/envs/assets_v2/sawyer_xyz/sawyer_assembly_peg.xml"
):
model_path = "/home/mdalal/research/metaworld/metaworld/envs/assets_v2/sawyer_xyz/sawyer_assembly_peg_torque.xml"
if type(self) == SawyerNutAssemblyEnvV2:
self.reset_qpos = [
1.88500731e00,
-5.88898933e-01,
-9.64183847e-01,
1.64749509e00,
9.28632075e-01,
1.02969815e00,
2.31183042e00,
-1.71909704e-04,
1.71743732e-04,
0.00000000e00,
6.00000024e-01,
1.99999996e-02,
7.07106763e-01,
1.36367940e-04,
1.60057621e-04,
7.07106768e-01,
]
else:
self.reset_qpos = [
2.16526293e00,
-5.57054189e-01,
-1.22832464e00,
2.17377367e00,
1.24233511e00,
1.00706272e00,
1.91061865e00,
-1.71855687e-04,
1.71795647e-04,
6.25459890e-02,
7.42607147e-01,
2.50073207e-02,
7.07106763e-01,
1.36367940e-04,
1.60057621e-04,
7.07106768e-01,
]
else:
raise ValueError("no model found")
self.frame_skip = frame_skip
self._use_dm_backend = True
# old zoomed out view
# camera_settings = {
# "distance": 1.878359835328275,
# "lookat": [0.16854934, 0.27084485, 0.23161897],
# "azimuth": 141.328125,
# "elevation": -53.203125160653144,
# }
# intermediate view
# camera_settings = {
# "distance":0.8304056576521722,
# "lookat":[0.21052547, 0.32329237, 0.587819 ],
# "azimuth": 141.328125,
# "elevation": -53.203125160653144,
# }
# super zoomed in - working
# camera_settings = {
# "distance": 0.38227044687537043,
# "lookat": [0.21052547, 0.32329237, 0.587819],
# "azimuth": 141.328125,
# "elevation": -53.203125160653144,
# }
# side view - semi-close up
# camera_settings = {
# "distance":0.513599996134662,
# "lookat":[0.28850459, 0.56757972, 0.54530015],
# "azimuth": 178.9453125,
# "elevation": -60.00000040512532,
# }
# side view super zoomed in
camera_settings = {
"distance": 0.31785821791481395,
"lookat": [0.28850459, 0.56757972, 0.54530015],
"azimuth": 178.59375,
"elevation": -60.46875041909516,
}
if self._use_dm_backend:
dm_mujoco = module.get_dm_mujoco()
if model_path.endswith(".mjb"):
self.sim = dm_mujoco.Physics.from_binary_path(model_path)
else:
self.sim = dm_mujoco.Physics.from_xml_path(model_path)
self.model = self.sim.model
patch_mjlib_accessors(self.get_mjlib(), self.model, self.sim.data)
self.renderer = DMRenderer(self.sim,
camera_settings=camera_settings)
else: # Use mujoco_py
mujoco_py = module.get_mujoco_py()
self.model = mujoco_py.load_model_from_path(model_path)
self.sim = mujoco_py.MjSim(self.model)
self.renderer = MjPyRenderer(self.sim,
camera_settings=camera_settings)
self.data = self.sim.data
self.viewer = None
self._viewers = {}
self._rgb_array_res = rgb_array_res
self.metadata = {
"render.modes": ["human"],
"video.frames_per_second": int(np.round(1.0 / self.dt)),
}
self.init_qpos = self.sim.data.qpos.ravel().copy()
self.init_qvel = self.sim.data.qvel.ravel().copy()
self._did_see_sim_exception = False
self.seed()