def test_rendering():
model = load_model_from_xml(BASIC_MODEL_XML)
sim = MjSim(model)
sim.forward()
img, depth = sim.render(200, 200, depth=True)
assert img.shape == (200, 200, 3)
compare_imgs(img, 'test_rendering.freecam.png')
depth = (depth - np.min(depth)) / (np.max(depth) - np.min(depth))
depth = np.asarray(depth * 255, dtype=np.uint8)
assert depth.shape == (200, 200)
compare_imgs(depth, 'test_rendering.freecam.depth.png')
img = sim.render(100, 100, camera_name="camera1")
assert img.shape == (100, 100, 3)
compare_imgs(img, 'test_rendering.camera1.png')
img = sim.render(200, 100, camera_name="camera1")
assert img.shape == (100, 200, 3)
compare_imgs(img, 'test_rendering.camera1.narrow.png')
render_context = sim.render_contexts[0]
render_context.add_marker(size=np.array([.4, .5, .6]),
pos=np.array([.4, .5, .6]),
rgba=np.array([.7, .8, .9, 1.0]),
label="mark")
img = sim.render(200, 200, camera_name="camera1")
assert img.shape == (200, 200, 3)
compare_imgs(img, 'test_rendering_markers.camera1.png')
def find_contact_height(sim: mj.MjSim, iterations: int = 10) -> float:
""" Find the height we should put our walker at so we are just barely in contact with the ground.
Runs a binary search on height (assumed to be state[1][2]) such that the body with the minimum z coordinate
on the ground is just barely in contact with the ground.
Args:
sim: the instantiated sim MjSim object you want to find the correct height for
iterations: number of times to run_util the binary search
Returns:
just the height as a float, you'll need to set it yourself
"""
state = sim.get_state()
height_guess = state[1][1]
step = height_guess * 2 # upper range for how much we will search for
for _ in range(iterations):
if sim.data.ncon: # ncon is the number of collisions
height_guess += step
else:
height_guess -= step
state[1][1] = height_guess
sim.set_state(state)
sim.forward()
step /= 2
return height_guess
def test_rendering():
model = load_model_from_xml(BASIC_MODEL_XML)
sim = MjSim(model)
sim.forward()
img, depth = sim.render(200, 200, depth=True)
assert img.shape == (200, 200, 3)
compare_imgs(img, 'test_rendering.freecam.png')
depth = (depth - np.min(depth)) / (np.max(depth) - np.min(depth))
depth = np.asarray(depth * 255, dtype=np.uint8)
assert depth.shape == (200, 200)
# Unfortunately mujoco 2.0 renders slightly different depth image on mac and on linux here
if "darwin" in sys.platform.lower():
compare_imgs(depth, 'test_rendering.freecam.depth-darwin.png')
else:
compare_imgs(depth, 'test_rendering.freecam.depth.png')
img = sim.render(100, 100, camera_name="camera1")
assert img.shape == (100, 100, 3)
compare_imgs(img, 'test_rendering.camera1.png')
img = sim.render(200, 100, camera_name="camera1")
assert img.shape == (100, 200, 3)
compare_imgs(img, 'test_rendering.camera1.narrow.png')
render_context = sim.render_contexts[0]
render_context.add_marker(size=np.array([.4, .5, .6]),
pos=np.array([.4, .5, .6]),
rgba=np.array([.7, .8, .9, 1.0]),
label="mark")
img = sim.render(200, 200, camera_name="camera1")
assert img.shape == (200, 200, 3)
compare_imgs(img, 'test_rendering_markers.camera1.png')
def test_multiple_sims():
# Ensure that creating new simulators still produces good renderings.
xml = """
<mujoco>
<asset>
<texture name="t1" width="32" height="32" type="2d" builtin="flat" />
<material name="m1" texture="t1" />
</asset>
<worldbody>
<light diffuse=".5 .5 .5" pos="0 0 5" dir="0 0 -1" />
<camera name="topcam" pos="0 0 2.5" zaxis="0 0 1" />
<geom name="g1" pos="0 0 0" type="box" size="1 1 0.1" rgba="1 1 1 1" material="m1" />
</worldbody>
</mujoco>
"""
model = load_model_from_xml(xml)
random_state = np.random.RandomState(0)
for i in range(3):
sim = MjSim(model)
sim.forward()
modder = TextureModder(sim, random_state=random_state)
for j in range(2):
modder.rand_checker('g1')
compare_imgs(sim.render(201, 205, camera_name="topcam"),
'test_multiple_sims.loop%d_%d.png' % (i, j))
def test_multiple_sims():
# Ensure that creating new simulators still produces good renderings.
xml = """
<mujoco>
<asset>
<texture name="t1" width="32" height="32" type="2d" builtin="flat" />
<material name="m1" texture="t1" />
</asset>
<worldbody>
<light diffuse=".5 .5 .5" pos="0 0 5" dir="0 0 -1" />
<camera name="topcam" pos="0 0 2.5" zaxis="0 0 1" />
<geom name="g1" pos="0 0 0" type="box" size="1 1 0.1" rgba="1 1 1 1" material="m1" />
</worldbody>
</mujoco>
"""
model = load_model_from_xml(xml)
random_state = np.random.RandomState(0)
for i in range(3):
sim = MjSim(model)
sim.forward()
modder = TextureModder(sim, random_state=random_state)
for j in range(2):
modder.rand_checker('g1')
compare_imgs(
sim.render(201, 205, camera_name="topcam"),
'test_multiple_sims.loop%d_%d.png' % (i, j))
def main():
parent_dir_path = str(pathlib.Path(__file__).parent.absolute())
fname = parent_dir_path + '/kinova_j2s6s300/j2s6s300.xml'
model = load_model_from_path(fname)
sim = MjSim(model)
viewer = MjViewer(sim)
t = 0
while True:
if t == 1000:
ndof = len(sim.data.qpos)
captured_state = copy.deepcopy(sim.data.qpos)
desired_vel = [0] * len(captured_state)
kv = np.eye(ndof) * 10
if t < 1000:
sim.data.ctrl[:] = sim.data.qfrc_bias[:]
else:
sim.data.ctrl[:] = mjc.pd(None,
desired_vel,
captured_state,
sim,
ndof=len(captured_state),
kv=kv)
sim.forward()
sim.step()
viewer.render()
t += 1
def test_rendering():
model = load_model_from_xml(BASIC_MODEL_XML)
sim = MjSim(model)
sim.forward()
img, depth = sim.render(200, 200, depth=True)
assert img.shape == (200, 200, 3)
compare_imgs(img, 'test_rendering.freecam.png')
depth = (depth - np.min(depth)) / (np.max(depth) - np.min(depth))
depth = np.asarray(depth * 255, dtype=np.uint8)
assert depth.shape == (200, 200)
compare_imgs(depth, 'test_rendering.freecam.depth.png')
img = sim.render(100, 100, camera_name="camera1")
assert img.shape == (100, 100, 3)
compare_imgs(img, 'test_rendering.camera1.png')
img = sim.render(200, 100, camera_name="camera1")
assert img.shape == (100, 200, 3)
compare_imgs(img, 'test_rendering.camera1.narrow.png')
render_context = sim.render_contexts[0]
render_context.add_marker(size=np.array([.4, .5, .6]),
pos=np.array([.4, .5, .6]),
rgba=np.array([.7, .8, .9, 1.0]),
label="mark")
img = sim.render(200, 200, camera_name="camera1")
assert img.shape == (200, 200, 3)
compare_imgs(img, 'test_rendering_markers.camera1.png')
def test_rendering():
model = load_model_from_xml(BASIC_MODEL_XML)
sim = MjSim(model)
sim.forward()
img, depth = sim.render(200, 200, depth=True)
assert img.shape == (200, 200, 3)
compare_imgs(img, 'test_rendering.freecam.png')
depth = (depth - np.min(depth)) / (np.max(depth) - np.min(depth))
depth = np.asarray(depth * 255, dtype=np.uint8)
assert depth.shape == (200, 200)
# Unfortunately mujoco 2.0 renders slightly different depth image on mac and on linux here
if "darwin" in sys.platform.lower():
compare_imgs(depth, 'test_rendering.freecam.depth-darwin.png')
else:
compare_imgs(depth, 'test_rendering.freecam.depth.png')
img = sim.render(100, 100, camera_name="camera1")
assert img.shape == (100, 100, 3)
compare_imgs(img, 'test_rendering.camera1.png')
img = sim.render(200, 100, camera_name="camera1")
assert img.shape == (100, 200, 3)
compare_imgs(img, 'test_rendering.camera1.narrow.png')
render_context = sim.render_contexts[0]
render_context.add_marker(size=np.array([.4, .5, .6]),
pos=np.array([.4, .5, .6]),
rgba=np.array([.7, .8, .9, 1.0]),
label="mark")
img = sim.render(200, 200, camera_name="camera1")
assert img.shape == (200, 200, 3)
compare_imgs(img, 'test_rendering_markers.camera1.png')
def test_data_attribute_getters():
model = load_model_from_xml(BASIC_MODEL_XML)
sim = MjSim(model)
sim.forward()
assert_array_equal(sim.data.get_body_xpos("body1"), [0, 0, 1])
with pytest.raises(ValueError):
sim.data.get_body_xpos("body_foo")
with pytest.raises(RuntimeError):
sim.data.get_xpos("body1")
assert len(sim.data.get_body_xquat("body1")) == 4
assert_array_equal(sim.data.get_body_xmat("body1").shape, (3, 3))
# At (0, 1, 1) since the geom is displaced in the body
assert_array_equal(sim.data.get_body_xipos("body1"), [0, 1, 1])
assert_array_equal(sim.data.get_site_xpos("site1"), [1, 0, 1])
assert_array_equal(sim.data.get_site_xmat("site1").shape, (3, 3))
assert_array_equal(sim.data.get_geom_xpos("geom1"), [0.5, 0.4, 0.3])
assert_array_equal(sim.data.get_geom_xpos("geom2"), [0, 1, 1])
assert_array_equal(sim.data.get_geom_xmat("geom2").shape, (3, 3))
assert_array_equal(sim.data.get_light_xpos("light1"), [0, 0, 3])
assert_array_equal(sim.data.get_light_xdir("light1"), [0, 0, -1])
assert_array_equal(sim.data.get_camera_xpos("camera1"), [3, 0, 0])
assert_array_equal(sim.data.get_camera_xmat("camera1").shape, (3, 3))
assert_array_equal(sim.data.get_joint_xaxis("joint1"), [0, 0, 1])
assert_array_equal(sim.data.get_joint_xanchor("joint1"), [0, 0, 1])
def test_arrays_of_objs():
model = load_model_from_xml(BASIC_MODEL_XML)
sim = MjSim(model)
sim.forward()
renderer = cymj.MjRenderContext(sim, offscreen=True)
assert len(
renderer.scn.camera) == 2, "Expecting scn.camera to be available"
def test_data_attribute_getters():
model = load_model_from_xml(BASIC_MODEL_XML)
sim = MjSim(model)
sim.forward()
assert_array_equal(sim.data.get_body_xpos("body1"), [0, 0, 1])
with pytest.raises(ValueError):
sim.data.get_body_xpos("body_foo")
with pytest.raises(RuntimeError):
sim.data.get_xpos("body1")
assert len(sim.data.get_body_xquat("body1")) == 4
assert_array_equal(sim.data.get_body_xmat("body1").shape, (3, 3))
# At (0, 1, 1) since the geom is displaced in the body
assert_array_equal(sim.data.get_body_xipos("body1"), [0, 1, 1])
assert_array_equal(sim.data.get_site_xpos("site1"), [1, 0, 1])
assert_array_equal(sim.data.get_site_xmat("site1").shape, (3, 3))
assert_array_equal(sim.data.get_geom_xpos("geom1"), [0.5, 0.4, 0.3])
assert_array_equal(sim.data.get_geom_xpos("geom2"), [0, 1, 1])
assert_array_equal(sim.data.get_geom_xmat("geom2").shape, (3, 3))
assert_array_equal(sim.data.get_light_xpos("light1"), [0, 0, 3])
assert_array_equal(sim.data.get_light_xdir("light1"), [0, 0, -1])
assert_array_equal(sim.data.get_camera_xpos("camera1"), [3, 0, 0])
assert_array_equal(sim.data.get_camera_xmat("camera1").shape, (3, 3))
assert_array_equal(sim.data.get_joint_xaxis("joint1"), [0, 0, 1])
assert_array_equal(sim.data.get_joint_xanchor("joint1"), [0, 0, 1])
def play(self, mocap_filepath):
from mujoco_py import load_model_from_xml, MjSim, MjViewer
curr_path = getcwd()
xmlpath = '/mujoco/humanoid_deepmimic/envs/asset/dp_env_v2.xml'
with open(curr_path + xmlpath) as fin:
MODEL_XML = fin.read()
model = load_model_from_xml(MODEL_XML)
sim = MjSim(model)
viewer = MjViewer(sim)
self.read_raw_data(mocap_filepath)
self.convert_raw_data()
from time import sleep
phase_offset = np.array([0.0, 0.0, 0.0])
while True:
for k in range(len(self.data)):
tmp_val = self.data_config[k]
sim_state = sim.get_state()
sim_state.qpos[:] = tmp_val[:]
sim_state.qpos[:3] += phase_offset[:]
sim.set_state(sim_state)
sim.forward()
viewer.render()
sim_state = sim.get_state()
phase_offset = sim_state.qpos[:3]
phase_offset[2] = 0
def test_high_res():
model = load_model_from_xml(BASIC_MODEL_XML)
sim = MjSim(model)
sim.forward()
img = sim.render(1000, 1000)
img = scipy.misc.imresize(img, (200, 200, 3))
assert img.shape == (200, 200, 3)
compare_imgs(img, 'test_rendering.freecam.png')
def test_mj_sim_basics():
model = load_model_from_xml(BASIC_MODEL_XML)
sim = MjSim(model, nsubsteps=2)
sim.reset()
sim.step()
sim.reset()
sim.forward()
def test_high_res():
model = load_model_from_xml(BASIC_MODEL_XML)
sim = MjSim(model)
sim.forward()
img = sim.render(1000, 1000)
img = np.array(Image.fromarray(img).resize(size=(200, 200)))
assert img.shape == (200, 200, 3)
compare_imgs(img, 'test_rendering.freecam.png')
def test_high_res():
model = load_model_from_xml(BASIC_MODEL_XML)
sim = MjSim(model)
sim.forward()
img = sim.render(1000, 1000)
img = np.array(Image.fromarray(img).resize(size=(200, 200)))
assert img.shape == (200, 200, 3)
compare_imgs(img, 'test_rendering.freecam.png')
def test_high_res():
model = load_model_from_xml(BASIC_MODEL_XML)
sim = MjSim(model)
sim.forward()
img = sim.render(1000, 1000)
img = scipy.misc.imresize(img, (200, 200, 3))
assert img.shape == (200, 200, 3)
compare_imgs(img, 'test_rendering.freecam.png')
def test_mj_sim_basics():
model = load_model_from_xml(BASIC_MODEL_XML)
sim = MjSim(model, nsubsteps=2)
sim.reset()
sim.step()
sim.reset()
sim.forward()
def test_viewercontext():
model = load_model_from_xml(BASIC_MODEL_XML)
sim = MjSim(model)
sim.forward()
renderer = cymj.MjRenderContext(sim, offscreen=True)
renderer.add_marker(type=const.GEOM_SPHERE,
size=np.ones(3) * 0.1,
pos=np.zeros(3),
mat=np.eye(3).flatten(),
rgba=np.ones(4),
label="mark")
def test_viewercontext():
model = load_model_from_xml(BASIC_MODEL_XML)
sim = MjSim(model)
sim.forward()
renderer = cymj.MjRenderContext(sim, offscreen=True)
renderer.add_marker(type=const.GEOM_SPHERE,
size=np.ones(3) * 0.1,
pos=np.zeros(3),
mat=np.eye(3).flatten(),
rgba=np.ones(4),
label="mark")
class Robot:
''' Simple utility class for getting mujoco-specific info about a robot '''
def __init__(self, path):
base_path = os.path.join(BASE_DIR, path)
self.sim = MjSim(load_model_from_path(base_path))
self.sim.forward()
# Needed to figure out z-height of free joint of offset body
self.z_height = self.sim.data.get_body_xpos('robot')[2]
# Get a list of geoms in the robot
self.geom_names = [
n for n in self.sim.model.geom_names if n != 'floor'
]
# Needed to figure out the observation spaces
self.nq = self.sim.model.nq
self.nv = self.sim.model.nv
# Needed to figure out action space
self.nu = self.sim.model.nu
# Needed to figure out observation space
# See engine.py for an explanation for why we treat these separately
self.hinge_pos_names = []
self.hinge_vel_names = []
self.ballquat_names = []
self.ballangvel_names = []
self.sensor_dim = {}
for name in self.sim.model.sensor_names:
id = self.sim.model.sensor_name2id(name)
self.sensor_dim[name] = self.sim.model.sensor_dim[id]
sensor_type = self.sim.model.sensor_type[id]
if self.sim.model.sensor_objtype[id] == const.OBJ_JOINT:
joint_id = self.sim.model.sensor_objid[id]
joint_type = self.sim.model.jnt_type[joint_id]
if joint_type == const.JNT_HINGE:
if sensor_type == const.SENS_JOINTPOS:
self.hinge_pos_names.append(name)
elif sensor_type == const.SENS_JOINTVEL:
self.hinge_vel_names.append(name)
else:
t = self.sim.model.sensor_type[i]
raise ValueError(
'Unrecognized sensor type {} for joint'.format(t))
elif joint_type == const.JNT_BALL:
if sensor_type == const.SENS_BALLQUAT:
self.ballquat_names.append(name)
elif sensor_type == const.SENS_BALLANGVEL:
self.ballangvel_names.append(name)
elif joint_type == const.JNT_SLIDE:
# Adding slide joints is trivially easy in code,
# but this removes one of the good properties about our observations.
# (That we are invariant to relative whole-world transforms)
# If slide joints are added we sould ensure this stays true!
raise ValueError(
'Slide joints in robots not currently supported')
def test_read_depth_buffer():
model = load_model_from_xml(BASIC_MODEL_XML)
sim = MjSim(model)
sim.forward()
ctx = MjRenderContext(sim, offscreen=True, opengl_backend='glfw')
buf = np.zeros((11, 100), dtype=np.float32)
assert buf.sum() == 0, f'{buf.sum()}'
ctx.render(buf.shape[1], buf.shape[0], 0)
ctx.read_pixels_depth(buf)
assert buf.sum() != 0, f'{buf.sum()} {buf.max()}'
def test_textures():
model = load_model_from_xml(BASIC_MODEL_XML)
sim = MjSim(model)
sim.forward()
compare_imgs(sim.render(201, 205, camera_name="topcam"),
'test_textures.premod.png')
random_state = np.random.RandomState(0)
modder = TextureModder(sim, random_state=random_state)
modder.whiten_materials()
modder.whiten_materials(['g1', 'g2'])
modder.set_rgb('g1', (255, 0, 0))
modder.set_rgb('g2', (0, 255, 0))
modder.set_rgb('g3', (0, 0, 255))
modder.set_rgb('g4', (255, 0, 255))
compare_imgs(sim.render(201, 205, camera_name="topcam"),
'test_textures.rgb.png')
modder.set_checker('g1', (255, 0, 0), (0, 255, 0))
modder.set_gradient('g2', (0, 255, 0), (0, 0, 255), vertical=True)
modder.set_gradient('g3', (255, 255, 0), (0, 0, 255), vertical=False)
modder.set_noise('g4', (0, 0, 255), (255, 0, 0), 0.1)
compare_imgs(sim.render(201, 205, camera_name="topcam"),
'test_textures.variety.png')
modder.rand_checker('g1')
modder.rand_gradient('g2')
modder.rand_noise('g3')
modder.rand_rgb('g4')
compare_imgs(sim.render(201, 205, camera_name="topcam"),
'test_textures.rand_specific.png')
modder.rand_all('g1')
modder.rand_all('g2')
modder.rand_all('g3')
modder.rand_all('g4')
compare_imgs(sim.render(201, 205, camera_name="topcam"),
'test_textures.rand_all.png')
modder.rand_checker('g1')
modder.rand_checker('g2')
modder.rand_checker('g3')
modder.rand_checker('g4')
mat_modder = MaterialModder(sim, random_state=random_state)
mat_modder.rand_texrepeat('g1')
mat_modder.rand_texrepeat('g2')
mat_modder.rand_texrepeat('g3')
mat_modder.rand_texrepeat('g4')
compare_imgs(sim.render(201, 205, camera_name="topcam"),
'test_textures.rand_texrepeat.png')
def test_textures():
model = load_model_from_xml(BASIC_MODEL_XML)
sim = MjSim(model)
sim.forward()
compare_imgs(sim.render(201, 205, camera_name="topcam"),
'test_textures.premod.png')
random_state = np.random.RandomState(0)
modder = TextureModder(sim, random_state=random_state)
modder.whiten_materials()
modder.whiten_materials(['g1', 'g2'])
modder.set_rgb('g1', (255, 0, 0))
modder.set_rgb('g2', (0, 255, 0))
modder.set_rgb('g3', (0, 0, 255))
modder.set_rgb('g4', (255, 0, 255))
compare_imgs(sim.render(201, 205, camera_name="topcam"),
'test_textures.rgb.png')
modder.set_checker('g1', (255, 0, 0), (0, 255, 0))
modder.set_gradient('g2', (0, 255, 0), (0, 0, 255), vertical=True)
modder.set_gradient('g3', (255, 255, 0), (0, 0, 255), vertical=False)
modder.set_noise('g4', (0, 0, 255), (255, 0, 0), 0.1)
compare_imgs(sim.render(201, 205, camera_name="topcam"),
'test_textures.variety.png')
modder.rand_checker('g1')
modder.rand_gradient('g2')
modder.rand_noise('g3')
modder.rand_rgb('g4')
compare_imgs(sim.render(201, 205, camera_name="topcam"),
'test_textures.rand_specific.png')
modder.rand_all('g1')
modder.rand_all('g2')
modder.rand_all('g3')
modder.rand_all('g4')
compare_imgs(sim.render(201, 205, camera_name="topcam"),
'test_textures.rand_all.png')
modder.rand_checker('g1')
modder.rand_checker('g2')
modder.rand_checker('g3')
modder.rand_checker('g4')
mat_modder = MaterialModder(sim, random_state=random_state)
mat_modder.rand_texrepeat('g1')
mat_modder.rand_texrepeat('g2')
mat_modder.rand_texrepeat('g3')
mat_modder.rand_texrepeat('g4')
compare_imgs(sim.render(201, 205, camera_name="topcam"),
'test_textures.rand_texrepeat.png')
def test_glfw_context():
model = load_model_from_xml(BASIC_MODEL_XML)
sim = MjSim(model)
sim.forward()
render_context = MjRenderContext(sim, offscreen=True, opengl_backend='glfw')
assert len(sim.render_contexts) == 1
assert sim.render_contexts[0] is render_context
assert isinstance(render_context.opengl_context, GlfwContext)
compare_imgs(sim.render(201, 205, camera_name="topcam"), 'test_glfw_context.png')
assert len(sim.render_contexts) == 1
assert sim.render_contexts[0] is render_context
def test_glfw_context():
model = load_model_from_xml(BASIC_MODEL_XML)
sim = MjSim(model)
sim.forward()
render_context = MjRenderContext(sim, offscreen=True, opengl_backend='glfw')
assert len(sim.render_contexts) == 1
assert sim.render_contexts[0] is render_context
assert isinstance(render_context.opengl_context, GlfwContext)
compare_imgs(sim.render(201, 205, camera_name="topcam"), 'test_glfw_context.png')
assert len(sim.render_contexts) == 1
assert sim.render_contexts[0] is render_context
def test_joint_qpos_qvel_ops():
model = load_model_from_xml(BASIC_MODEL_XML)
sim = MjSim(model)
sim.forward()
# Test setting one with a list
sim.data.set_joint_qpos("joint1", [1, 2, 3, 1, 0, 0, 0])
# And the other with an np.ndarray
sim.data.set_joint_qvel("joint1", np.array([1, 2, 3, 0.1, 0.1, 0.1]))
sim.forward()
assert_array_equal(sim.data.get_joint_qpos(
"joint1"), [1, 2, 3, 1, 0, 0, 0])
assert_array_equal(sim.data.get_joint_qvel(
"joint1"), [1, 2, 3, 0.1, 0.1, 0.1])
def test_joint_qpos_qvel_ops():
model = load_model_from_xml(BASIC_MODEL_XML)
sim = MjSim(model)
sim.forward()
# Test setting one with a list
sim.data.set_joint_qpos("joint1", [1, 2, 3, 1, 0, 0, 0])
# And the other with an np.ndarray
sim.data.set_joint_qvel("joint1", np.array([1, 2, 3, 0.1, 0.1, 0.1]))
sim.forward()
assert_array_equal(sim.data.get_joint_qpos(
"joint1"), [1, 2, 3, 1, 0, 0, 0])
assert_array_equal(sim.data.get_joint_qvel(
"joint1"), [1, 2, 3, 0.1, 0.1, 0.1])
def test_rendering_failing():
model = load_model_from_xml(BASIC_MODEL_XML)
sim = MjSim(model)
sim.forward()
sim.render(100, 100)
render_context = sim.render_contexts[0]
render_context.add_marker(size=np.array([.4, .5, .6]),
pos=np.array([.4, .5, .6]),
rgba=np.array([.7, .8, .9, 1.0]),
label="blaaaa")
img = sim.render(200, 200, camera_name="camera1")
assert img.shape == (200, 200, 3)
try:
compare_imgs(img, 'test_rendering_markers.camera1.png')
assert False
except Exception as e:
pass
def test_rendering_failing():
model = load_model_from_xml(BASIC_MODEL_XML)
sim = MjSim(model)
sim.forward()
sim.render(100, 100)
render_context = sim.render_contexts[0]
render_context.add_marker(size=np.array([.4, .5, .6]),
pos=np.array([.4, .5, .6]),
rgba=np.array([.7, .8, .9, 1.0]),
label="blaaaa")
img = sim.render(200, 200, camera_name="camera1")
assert img.shape == (200, 200, 3)
try:
compare_imgs(img, 'test_rendering_markers.camera1.png')
assert False
except Exception as e:
pass
def test_ray(self):
''' Test raycasting and exclusions '''
sim = MjSim(load_model_from_xml(self.xml))
sim.forward()
# Include all geoms
self.check_rays(sim,
[0.9, 0.1, 0.9, 0.1, 0.9, 0.1, -1.0],
['A', 'A', 'B', 'B', 'C', 'C', None])
# Include static geoms, but exclude worldbody (which contains 'A')
self.check_rays(sim,
[2.9, 1.9, 0.9, 0.1, 0.9, 0.1, -1.0],
['B', 'B', 'B', 'B', 'C', 'C', None],
exclude_body=0)
# Include static geoms, and exclude body 1 (which contains 'C')
self.check_rays(sim,
[0.9, 0.1, 0.9, 0.1, -1.0, -1.0, -1.0],
['A', 'A', 'B', 'B', None, None, None],
exclude_body=1)
# Include static geoms, and exclude body 2 (which contains 'B')
self.check_rays(sim,
[0.9, 0.1, 2.9, 1.9, 0.9, 0.1, -1.0],
['A', 'A', 'C', 'C', 'C', 'C', None],
exclude_body=2)
# Exclude static geoms ('A' is the only static geom)
self.check_rays(sim,
[2.9, 1.9, 0.9, 0.1, 0.9, 0.1, -1.0],
['B', 'B', 'B', 'B', 'C', 'C', None],
include_static_geoms=False)
# Exclude static geoms, and exclude body 1 ('C')
self.check_rays(sim,
[2.9, 1.9, 0.9, 0.1, -1.0, -1.0, -1.0],
['B', 'B', 'B', 'B', None, None, None],
include_static_geoms=False, exclude_body=1)
# Exclude static geoms, and exclude body 2 (which contains 'B')
self.check_rays(sim,
[4.9, 3.9, 2.9, 1.9, 0.9, 0.1, -1.0],
['C', 'C', 'C', 'C', 'C', 'C', None],
include_static_geoms=False, exclude_body=2)
def test_ray(self):
''' Test raycasting and exclusions '''
sim = MjSim(load_model_from_xml(self.xml))
sim.forward()
# Include all geoms
self.check_rays(sim,
[0.9, 0.1, 0.9, 0.1, 0.9, 0.1, -1.0],
['A', 'A', 'B', 'B', 'C', 'C', None])
# Include static geoms, but exclude worldbody (which contains 'A')
self.check_rays(sim,
[2.9, 1.9, 0.9, 0.1, 0.9, 0.1, -1.0],
['B', 'B', 'B', 'B', 'C', 'C', None],
exclude_body=0)
# Include static geoms, and exclude body 1 (which contains 'C')
self.check_rays(sim,
[0.9, 0.1, 0.9, 0.1, -1.0, -1.0, -1.0],
['A', 'A', 'B', 'B', None, None, None],
exclude_body=1)
# Include static geoms, and exclude body 2 (which contains 'B')
self.check_rays(sim,
[0.9, 0.1, 2.9, 1.9, 0.9, 0.1, -1.0],
['A', 'A', 'C', 'C', 'C', 'C', None],
exclude_body=2)
# Exclude static geoms ('A' is the only static geom)
self.check_rays(sim,
[2.9, 1.9, 0.9, 0.1, 0.9, 0.1, -1.0],
['B', 'B', 'B', 'B', 'C', 'C', None],
include_static_geoms=False)
# Exclude static geoms, and exclude body 1 ('C')
self.check_rays(sim,
[2.9, 1.9, 0.9, 0.1, -1.0, -1.0, -1.0],
['B', 'B', 'B', 'B', None, None, None],
include_static_geoms=False, exclude_body=1)
# Exclude static geoms, and exclude body 2 (which contains 'B')
self.check_rays(sim,
[4.9, 3.9, 2.9, 1.9, 0.9, 0.1, -1.0],
['C', 'C', 'C', 'C', 'C', 'C', None],
include_static_geoms=False, exclude_body=2)
class Mujoco_Renderer():
def __init__(self, mujoco_xml, hp):
from mujoco_py import load_model_from_path, MjSim
mujoco_xml = '/'.join(str.split(gcp.__file__, '/')[:-1]) + '/' + mujoco_xml
self.sim = MjSim(load_model_from_path(mujoco_xml))
self._hp = hp
def render(self, qpos):
sim_state = self.sim.get_state()
sim_state.qpos[:2] = qpos
sim_state.qvel[:] = np.zeros_like(self.sim.data.qvel)
self.sim.set_state(sim_state)
self.sim.forward()
subgoal_image = self.sim.render(self._hp.mpar.img_sz, self._hp.mpar.img_sz, camera_name='maincam')
# plt.imshow(subgoal_image)
# plt.savefig('test.png')
return subgoal_image
def test_mocap_ops():
model = load_model_from_xml(BASIC_MODEL_XML)
sim = MjSim(model)
sim.forward()
assert_array_equal(sim.data.get_body_xpos("mocap1"), [1, 0, 0])
assert_array_equal(sim.data.get_mocap_pos("mocap1"), [1, 0, 0])
assert_array_equal(sim.data.get_mocap_quat("mocap1"), [1, 0, 0, 0])
new_pos = [2, 1, 1]
new_quat = [0.707107, 0.707107, 0, 0]
sim.data.set_mocap_pos("mocap1", new_pos)
sim.data.set_mocap_quat("mocap1", new_quat)
sim.forward()
assert_array_equal(sim.data.get_mocap_pos("mocap1"), new_pos)
assert_array_almost_equal(sim.data.get_mocap_quat("mocap1"), new_quat)
assert_array_equal(sim.data.get_body_xpos("mocap1"), new_pos)
assert_array_almost_equal(sim.data.get_body_xquat("mocap1"), new_quat)
assert_array_almost_equal(sim.data.get_body_xmat("mocap1"),
[[1, 0, 0], [0, 0, -1], [0, 1, 0]])
class MjJacoEnv(object):
"""docstring for MjJacoEnv."""
def __init__(self, vis=False):
super(MjJacoEnv, self).__init__()
parent_dir_path = str(pathlib.Path(__file__).parent.absolute())
self.fname = parent_dir_path + '/jaco/jaco.xml'
self.model = load_model_from_path(self.fname)
self.sim = MjSim(self.model)
self.viewer = MjViewer(self.sim)
self.vis = vis
def step(self, action):
for i in range(len(action)):
self.sim.data.ctrl[i] = action[i]
self.sim.forward()
self.sim.step()
self.viewer.render() if self.vis else None
return self.sim.data.qpos
def test_mocap_ops():
model = load_model_from_xml(BASIC_MODEL_XML)
sim = MjSim(model)
sim.forward()
assert_array_equal(sim.data.get_body_xpos("mocap1"), [1, 0, 0])
assert_array_equal(sim.data.get_mocap_pos("mocap1"), [1, 0, 0])
assert_array_equal(sim.data.get_mocap_quat("mocap1"), [1, 0, 0, 0])
new_pos = [2, 1, 1]
new_quat = [0.707107, 0.707107, 0, 0]
sim.data.set_mocap_pos("mocap1", new_pos)
sim.data.set_mocap_quat("mocap1", new_quat)
sim.forward()
assert_array_equal(sim.data.get_mocap_pos("mocap1"), new_pos)
assert_array_almost_equal(sim.data.get_mocap_quat("mocap1"), new_quat)
assert_array_equal(sim.data.get_body_xpos("mocap1"), new_pos)
assert_array_almost_equal(sim.data.get_body_xquat("mocap1"), new_quat)
assert_array_almost_equal(sim.data.get_body_xmat("mocap1"),
[[1, 0, 0], [0, 0, -1], [0, 1, 0]])
def test_materials():
model = load_model_from_xml(BASIC_MODEL_XML)
sim = MjSim(model)
sim.forward()
compare_imgs(sim.render(201, 205, camera_name="topcam"),
'test_materials.premod.png')
random_state = np.random.RandomState(0)
modder = MaterialModder(sim, random_state=random_state)
modder.set_specularity('g1', 1.0)
modder.set_reflectance('g2', 1.0)
modder.set_shininess('g3', 1.0)
compare_imgs(sim.render(201, 205, camera_name="topcam"),
'test_materials.props.png')
modder.rand_all('g4')
compare_imgs(sim.render(201, 205, camera_name="topcam"),
'test_materials.rand_all.png')
def test_materials():
model = load_model_from_xml(BASIC_MODEL_XML)
sim = MjSim(model)
sim.forward()
compare_imgs(sim.render(201, 205, camera_name="topcam"),
'test_materials.premod.png')
random_state = np.random.RandomState(0)
modder = MaterialModder(sim, random_state=random_state)
modder.set_specularity('g1', 1.0)
modder.set_reflectance('g2', 1.0)
modder.set_shininess('g3', 1.0)
compare_imgs(sim.render(201, 205, camera_name="topcam"),
'test_materials.props.png')
modder.rand_all('g4')
compare_imgs(sim.render(201, 205, camera_name="topcam"),
'test_materials.rand_all.png')
def test_xvelr(): # xvelr = rotational velocity in world frame
xml = """
<mujoco>
<worldbody>
<body name="body1" pos="0 0 0">
<joint name="a" axis="1 0 0" pos="0 0 0" type="hinge"/>
<geom name="geom1" pos="0 0 0" size="0.3"/>
<body name="body2" pos="0 0 1">
<joint name="b" axis="1 0 0" pos="0 0 0" type="hinge"/>
<geom name="geom2" pos="0 0 0" size="0.3"/>
<site name="site1" size="0.1"/>
</body>
</body>
</worldbody>
<actuator>
<motor joint="a"/>
<motor joint="b"/>
</actuator>
</mujoco>
"""
model = load_model_from_xml(xml)
sim = MjSim(model)
sim.reset()
sim.forward()
# Check that xvelr starts out at zero (since qvel is zero)
site1_xvelr = sim.data.get_site_xvelr('site1')
np.testing.assert_allclose(site1_xvelr, np.zeros(3))
# Push the base body and step forward to get it moving
sim.data.ctrl[0] = 1e9
sim.step()
sim.forward()
# Check that the first body has nonzero xvelr
body1_xvelr = sim.data.get_body_xvelr('body1')
assert not np.allclose(body1_xvelr, np.zeros(3))
# Check that the second body has zero xvelr (still)
body2_xvelr = sim.data.get_body_xvelr('body2')
np.testing.assert_allclose(body2_xvelr, np.zeros(3))
# Check that this matches the batch (gathered) getter property
np.testing.assert_allclose(body2_xvelr, sim.data.body_xvelr[2])
def test_xvelr(): # xvelr = rotational velocity in world frame
xml = """
<mujoco>
<worldbody>
<body name="body1" pos="0 0 0">
<joint name="a" axis="1 0 0" pos="0 0 0" type="hinge"/>
<geom name="geom1" pos="0 0 0" size="0.3"/>
<body name="body2" pos="0 0 1">
<joint name="b" axis="1 0 0" pos="0 0 0" type="hinge"/>
<geom name="geom2" pos="0 0 0" size="0.3"/>
<site name="site1" size="0.1"/>
</body>
</body>
</worldbody>
<actuator>
<motor joint="a"/>
<motor joint="b"/>
</actuator>
</mujoco>
"""
model = load_model_from_xml(xml)
sim = MjSim(model)
sim.reset()
sim.forward()
# Check that xvelr starts out at zero (since qvel is zero)
site1_xvelr = sim.data.get_site_xvelr('site1')
np.testing.assert_allclose(site1_xvelr, np.zeros(3))
# Push the base body and step forward to get it moving
sim.data.ctrl[0] = 1e9
sim.step()
sim.forward()
# Check that the first body has nonzero xvelr
body1_xvelr = sim.data.get_body_xvelr('body1')
assert not np.allclose(body1_xvelr, np.zeros(3))
# Check that the second body has zero xvelr (still)
body2_xvelr = sim.data.get_body_xvelr('body2')
np.testing.assert_allclose(body2_xvelr, np.zeros(3))
# Check that this matches the batch (gathered) getter property
np.testing.assert_allclose(body2_xvelr, sim.data.body_xvelr[2])
class JuggleEnv:
def __init__(self):
self.control_freq: float = 50.0
self.control_timestep: float = 1.0 / self.control_freq
self.viewer = None
self.horizon = 1000
self.target = np.array([0.8, 0.0, 1.9])
# load model
self.robot: Robot = None
self.arena: Arena = None
self.pingpong: MujocoGeneratedObject = None
self.model: MujocoWorldBase = None
self._load_model()
# initialize simulation
self.mjpy_model = None
self.sim: MjSim = None
self.model_timestep: float = 0.0
self._initialize_sim()
# reset robot, object and internel variables
self.cur_time: float = 0.0
self.timestep: int = 0.0
self.done: bool = False
self._pingpong_body_id: int = -1
self._paddle_body_id: int = -1
self._reset_internel()
# internel variable for scoring
self._below_plane = False
self.plane_height = 1.5
def _load_model(self):
# Load the desired controller's default config as a dict
controller_config = load_controller_config(
default_controller="JOINT_VELOCITY")
controller_config["output_max"] = 1.0
controller_config["output_min"] = -1.0
robot_noise = {"magnitude": [0.05] * 7, "type": "gaussian"}
self.robot = SingleArm(
robot_type="IIWA",
idn=0,
controller_config=controller_config,
initial_qpos=[0.0, 0.7, 0.0, -1.4, 0.0, -0.56, 0.0],
initialization_noise=robot_noise,
gripper_type="PaddleGripper",
gripper_visualization=True,
control_freq=self.control_freq)
self.robot.load_model()
self.robot.robot_model.set_base_xpos([0, 0, 0])
self.arena = EmptyArena()
self.arena.set_origin([0.8, 0, 0])
self.pingpong = BallObject(name="pingpong",
size=[0.02],
rgba=[0.8, 0.8, 0, 1],
solref=[0.1, 0.03],
solimp=[0, 0, 1],
density=100)
pingpong_model = self.pingpong.get_collision()
pingpong_model.append(
new_joint(name="pingpong_free_joint", type="free"))
pingpong_model.set("pos", "0.8 0 2.0")
# merge into one
self.model = MujocoWorldBase()
self.model.merge(self.robot.robot_model)
self.model.merge(self.arena)
self.model.worldbody.append(pingpong_model)
def _initialize_sim(self):
# if we have an xml string, use that to create the sim. Otherwise, use the local model
self.mjpy_model = self.model.get_model(mode="mujoco_py")
# Create the simulation instance and run a single step to make sure changes have propagated through sim state
self.sim = MjSim(self.mjpy_model)
self.sim.step()
self.robot.reset_sim(self.sim)
self.model_timestep = self.sim.model.opt.timestep
def _reset_internel(self):
# reset robot
self.robot.setup_references()
self.robot.reset(deterministic=False)
# reset pingpong
pingpong_pos = self.target + np.random.rand(3) * 0.08 - 0.04
pingpong_quat = np.array([1.0, 0.0, 0.0, 0.0])
self.sim.data.set_joint_qpos(
"pingpong_free_joint",
np.concatenate([pingpong_pos, pingpong_quat]))
# get handle for important parts
self._pingpong_body_id = self.sim.model.body_name2id("pingpong")
self._paddle_body_id = self.sim.model.body_name2id(
"gripper0_paddle_body")
# Setup sim time based on control frequency
self.cur_time = 0
self.timestep = 0
self.done = False
def reset(self):
self.sim.reset()
self._reset_internel()
self.sim.forward()
return self._get_observation()
def _get_observation(self):
di = OrderedDict()
# get robot observation
di = self.robot.get_observations(di)
# get pingpong observation
pingpong_pos = np.array(
self.sim.data.body_xpos[self._pingpong_body_id])
di["pingpong_pos"] = pingpong_pos
return di
def step(self, action: np.ndarray):
if self.done:
raise ValueError("executing action in terminated episode")
policy_step = True
score = 0.0
for _ in range(int(self.control_timestep / self.model_timestep)):
self.sim.forward()
self.robot.control(action=action, policy_step=policy_step)
# self.sim.data.ctrl[:] = action*5.0
self.sim.step()
policy_step = False
# check if the ball pass the plane
h = self.sim.data.body_xpos[self._pingpong_body_id][2]
self._below_plane |= h < self.plane_height
if self._below_plane and h > self.plane_height:
score = 1.0
self._below_plane = False
self.timestep += 1
self.cur_time += self.control_timestep
observation = self._get_observation()
dist_xy = np.linalg.norm(
(observation["robot0_eef_pos"] - observation["pingpong_pos"])[:2])
# paddle_height = observation["robot0_eef_pos"][2]
self.done = self.timestep >= self.horizon or dist_xy > 0.2
reward = score # + 0 * (0.2 - dist_xy)
return observation, reward, self.done, {}
def render(self, mode="human"):
if mode == "human":
self._get_viewer().render()
elif mode == "rgb_array":
img = self.sim.render(1920, 1080)
return img[::-1, :, ::-1]
def _get_viewer(self):
if self.viewer is None:
self.viewer = MjViewer(self.sim)
self.viewer.vopt.geomgroup[0] = 0
self.viewer._hide_overlay = True
return self.viewer
def close(self):
self._destroy_viewer()
def _destroy_viewer(self):
if self.viewer is not None:
glfw.destroy_window(self.viewer.window)
self.viewer = None
def seed(self):
pass
class MujocoEnv(metaclass=EnvMeta):
"""
Initializes a Mujoco Environment.
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_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 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).
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
renderer (str): string for the renderer to use
renderer_config (dict): dictionary for the renderer configurations
Raises:
ValueError: [Invalid renderer selection]
"""
def __init__(
self,
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,
renderer="mujoco",
renderer_config=None,
):
# First, verify that both the on- and off-screen renderers are not being used simultaneously
if has_renderer is True and has_offscreen_renderer is True:
raise ValueError("the onscreen and offscreen renderers cannot be used simultaneously.")
# Rendering-specific attributes
self.has_renderer = has_renderer
self.has_offscreen_renderer = has_offscreen_renderer
self.render_camera = render_camera
self.render_collision_mesh = render_collision_mesh
self.render_visual_mesh = render_visual_mesh
self.render_gpu_device_id = render_gpu_device_id
self.viewer = None
# Simulation-specific attributes
self._observables = {} # Maps observable names to observable objects
self._obs_cache = {} # Maps observable names to pre-/partially-computed observable values
self.control_freq = control_freq
self.horizon = horizon
self.ignore_done = ignore_done
self.hard_reset = hard_reset
self._model_postprocessor = None # Function to post-process model after load_model() call
self.model = None
self.cur_time = None
self.model_timestep = None
self.control_timestep = None
self.deterministic_reset = False # Whether to add randomized resetting of objects / robot joints
self.renderer = renderer
self.renderer_config = renderer_config
# Load the model
self._load_model()
# Post-process model
self._postprocess_model()
# Initialize the simulation
self._initialize_sim()
#initializes the rendering
self.initialize_renderer()
# Run all further internal (re-)initialization required
self._reset_internal()
# Load observables
if hasattr(self.viewer, '_setup_observables'):
self._observables = self.viewer._setup_observables()
else:
self._observables = self._setup_observables()
# check if viewer has get observations method and set a flag for future use.
self.viewer_get_obs = hasattr(self.viewer, '_get_observations')
def initialize_renderer(self):
self.renderer = self.renderer.lower()
if self.renderer_config is None and self.renderer != 'mujoco':
self.renderer_config = load_renderer_config(self.renderer)
if self.renderer == 'mujoco' or self.renderer == 'default':
pass
elif self.renderer == 'nvisii':
from robosuite.renderers.nvisii.nvisii_renderer import NVISIIRenderer
self.viewer = NVISIIRenderer(env=self,
**self.renderer_config)
elif self.renderer == 'igibson':
from robosuite.renderers.igibson.igibson_renderer import iGibsonRenderer
self.viewer = iGibsonRenderer(env=self,
**self.renderer_config
)
else:
raise ValueError(f'{self.renderer} is not a valid renderer name. Valid options include default (native mujoco renderer), nvisii, and igibson')
def initialize_time(self, control_freq):
"""
Initializes the time constants used for simulation.
Args:
control_freq (float): Hz rate to run control loop at within the simulation
"""
self.cur_time = 0
self.model_timestep = macros.SIMULATION_TIMESTEP
if self.model_timestep <= 0:
raise ValueError("Invalid simulation timestep defined!")
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 set_model_postprocessor(self, postprocessor):
"""
Sets the post-processor function that self.model will be passed to after load_model() is called during resets.
Args:
postprocessor (None or function): If set, postprocessing method should take in a Task-based instance and
return no arguments.
"""
self._model_postprocessor = postprocessor
def _load_model(self):
"""Loads an xml model, puts it in self.model"""
pass
def _postprocess_model(self):
"""
Post-processes model after load_model() call. Useful for external objects (e.g.: wrappers) to
be able to modify the sim model before it is actually loaded into the simulation
"""
if self._model_postprocessor is not None:
self._model_postprocessor(self.model)
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.
"""
# Setup mappings from model to IDs
self.model.generate_id_mappings(sim=self.sim)
def _setup_observables(self):
"""
Sets up observables to be used for this environment.
Returns:
OrderedDict: Dictionary mapping observable names to its corresponding Observable object
"""
return OrderedDict()
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
"""
# if we have an xml string, use that to create the sim. Otherwise, use the local model
self.mjpy_model = load_model_from_xml(xml_string) if xml_string else self.model.get_model(mode="mujoco_py")
# Create the simulation instance and run a single step to make sure changes have propagated through sim state
self.sim = MjSim(self.mjpy_model)
self.sim.forward()
# Setup sim time based on control frequency
self.initialize_time(self.control_freq)
def reset(self):
"""
Resets simulation.
Returns:
OrderedDict: Environment observation space after reset occurs
"""
# TODO(yukez): investigate black screen of death
# Use hard reset if requested
if self.hard_reset and not self.deterministic_reset:
if self.renderer == 'mujoco' or self.renderer == 'default':
self._destroy_viewer()
self._load_model()
self._postprocess_model()
self._initialize_sim()
# Else, we only reset the sim internally
else:
self.sim.reset()
# Reset necessary robosuite-centric variables
self._reset_internal()
self.sim.forward()
# Setup observables, reloading if
self._obs_cache = {}
if self.hard_reset:
# If we're using hard reset, must re-update sensor object references
_observables = self._setup_observables()
for obs_name, obs in _observables.items():
self.modify_observable(observable_name=obs_name, attribute="sensor", modifier=obs._sensor)
# Make sure that all sites are toggled OFF by default
self.visualize(vis_settings={vis: False for vis in self._visualizations})
if self.viewer is not None and self.renderer != 'mujoco':
self.viewer.reset()
observations = self.viewer._get_observations(force_update=True) if self.viewer_get_obs else self._get_observations(force_update=True)
# Return new observations
return observations
def _reset_internal(self):
"""Resets simulation internal configurations."""
# create visualization screen or renderer
if self.has_renderer and self.viewer is None:
self.viewer = MujocoPyRenderer(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
# make sure mujoco-py doesn't block rendering frames
# (see https://github.com/StanfordVL/robosuite/issues/39)
self.viewer.viewer._render_every_frame = True
# Set the camera angle for viewing
if self.render_camera is not None:
self.viewer.set_camera(camera_id=self.sim.model.camera_name2id(self.render_camera))
elif self.has_offscreen_renderer:
if self.sim._render_context_offscreen is None:
render_context = MjRenderContextOffscreen(self.sim, device_id=self.render_gpu_device_id)
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._setup_references()
self.cur_time = 0
self.timestep = 0
self.done = False
# Empty observation cache and reset all observables
self._obs_cache = {}
for observable in self._observables.values():
observable.reset()
def _update_observables(self, force=False):
"""
Updates all observables in this environment
Args:
force (bool): If True, will force all the observables to update their internal values to the newest
value. This is useful if, e.g., you want to grab observations when directly setting simulation states
without actually stepping the simulation.
"""
for observable in self._observables.values():
observable.update(timestep=self.model_timestep, obs_cache=self._obs_cache, force=force)
def _get_observations(self, force_update=False):
"""
Grabs observations from the environment.
Args:
force_update (bool): If True, will force all the observables to update their internal values to the newest
value. This is useful if, e.g., you want to grab observations when directly setting simulation states
without actually stepping the simulation.
Returns:
OrderedDict: OrderedDict containing observations [(name_string, np.array), ...]
"""
observations = OrderedDict()
obs_by_modality = OrderedDict()
# Force an update if requested
if force_update:
self._update_observables(force=True)
# Loop through all observables and grab their current observation
for obs_name, observable in self._observables.items():
if observable.is_enabled() and observable.is_active():
obs = observable.obs
observations[obs_name] = obs
modality = observable.modality + "-state"
if modality not in obs_by_modality:
obs_by_modality[modality] = []
# Make sure all observations are numpy arrays so we can concatenate them
array_obs = [obs] if type(obs) in {int, float} or not obs.shape else obs
obs_by_modality[modality].append(np.array(array_obs))
# Add in modality observations
for modality, obs in obs_by_modality.items():
# To save memory, we only concatenate the image observations if explicitly requested
if modality == "image-state" and not macros.CONCATENATE_IMAGES:
continue
observations[modality] = np.concatenate(obs, axis=-1)
return observations
def step(self, action):
"""
Takes a step in simulation with control command @action.
Args:
action (np.array): Action to execute within the environment
Returns:
4-tuple:
- (OrderedDict) observations from the environment
- (float) reward from the environment
- (bool) whether the current episode is completed or not
- (dict) misc information
Raises:
ValueError: [Steps past episode termination]
"""
if self.done:
raise ValueError("executing action in terminated episode")
self.timestep += 1
# Since the env.step frequency is slower than the mjsim timestep frequency, the internal controller will output
# multiple torque commands in between new high level action commands. Therefore, we need to denote via
# 'policy_step' whether the current step we're taking is simply an internal update of the controller,
# or an actual policy update
policy_step = True
# Loop through the simulation at the model timestep rate until we're ready to take the next policy step
# (as defined by the control frequency specified at the environment level)
for i in range(int(self.control_timestep / self.model_timestep)):
self.sim.forward()
self._pre_action(action, policy_step)
self.sim.step()
self._update_observables()
policy_step = False
# Note: this is done all at once to avoid floating point inaccuracies
self.cur_time += self.control_timestep
reward, done, info = self._post_action(action)
if self.viewer is not None and self.renderer != 'mujoco':
self.viewer.update()
observations = self.viewer._get_observations() if self.viewer_get_obs else self._get_observations()
return observations, reward, done, info
def _pre_action(self, action, policy_step=False):
"""
Do any preprocessing before taking an action.
Args:
action (np.array): Action to execute within the environment
policy_step (bool): Whether this current loop is an actual policy step or internal sim update step
"""
self.sim.data.ctrl[:] = action
def _post_action(self, action):
"""
Do any housekeeping after taking an action.
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) empty dict to be filled with information by subclassed method
"""
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
Args:
action (np.array): Action to execute within the environment
Returns:
float: Reward from environment
"""
raise NotImplementedError
def render(self):
"""
Renders to an on-screen window.
"""
self.viewer.render()
def get_pixel_obs(self):
"""
Gets the pixel observations for the environment from the specified renderer
"""
self.viewer.get_pixel_obs()
def close_renderer(self):
"""
Closes the renderer
"""
self.viewer.close()
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.
Returns:
OrderedDict: Observations from the environment
"""
observation = self.viewer._get_observations() if self.viewer_get_obs else self._get_observations()
return observation
def clear_objects(self, object_names):
"""
Clears objects with the name @object_names out of the task space. This is useful
for supporting task modes with single types of objects, as in
@self.single_object_mode without changing the model definition.
Args:
object_names (str or list of str): Name of object(s) to remove from the task workspace
"""
object_names = {object_names} if type(object_names) is str else set(object_names)
for obj in self.model.mujoco_objects:
if obj.name in object_names:
self.sim.data.set_joint_qpos(obj.joints[0], np.array((10, 10, 10, 1, 0, 0, 0)))
def visualize(self, vis_settings):
"""
Do any needed visualization here
Args:
vis_settings (dict): Visualization keywords mapped to T/F, determining whether that specific
component should be visualized. Should have "env" keyword as well as any other relevant
options specified.
"""
# Set visuals for environment objects
for obj in self.model.mujoco_objects:
obj.set_sites_visibility(sim=self.sim, visible=vis_settings["env"])
def set_camera_pos_quat(self, camera_pos, camera_quat):
if self.renderer in ["nvisii", "igibson"]:
self.viewer.set_camera_pos_quat(camera_pos, camera_quat)
else:
raise AttributeError('setting camera position and quat requires renderer to be either NVISII or iGibson.')
def reset_from_xml_string(self, xml_string):
"""
Reloads the environment from an XML description of the environment.
Args:
xml_string (str): Filepath to the xml file that will be loaded directly into the sim
"""
# if there is an active viewer window, destroy it
if self.renderer != 'nvisii':
self.close()
# Since we are reloading from an xml_string, we are deterministically resetting
self.deterministic_reset = True
# initialize sim from xml
self._initialize_sim(xml_string=xml_string)
# Now reset as normal
self.reset()
# Turn off deterministic reset
self.deterministic_reset = False
def check_contact(self, geoms_1, geoms_2=None):
"""
Finds contact between two geom groups.
Args:
geoms_1 (str or list of str or MujocoModel): an individual geom name or list of geom names or a model. If
a MujocoModel is specified, the geoms checked will be its contact_geoms
geoms_2 (str or list of str or MujocoModel or None): another individual geom name or list of geom names.
If a MujocoModel is specified, the geoms checked will be its contact_geoms. If None, will check
any collision with @geoms_1 to any other geom in the environment
Returns:
bool: True if any geom in @geoms_1 is in contact with any geom in @geoms_2.
"""
# Check if either geoms_1 or geoms_2 is a string, convert to list if so
if type(geoms_1) is str:
geoms_1 = [geoms_1]
elif isinstance(geoms_1, MujocoModel):
geoms_1 = geoms_1.contact_geoms
if type(geoms_2) is str:
geoms_2 = [geoms_2]
elif isinstance(geoms_2, MujocoModel):
geoms_2 = geoms_2.contact_geoms
for contact in self.sim.data.contact[: 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 if geoms_2 is not None else True
# 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 geoms_2 is not None else True
if (c1_in_g1 and c2_in_g2) or (c1_in_g2 and c2_in_g1):
return True
return False
def get_contacts(self, model):
"""
Checks for any contacts with @model (as defined by @model's contact_geoms) and returns the set of
geom names 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 geoms that are actively in contact 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 and g2 not in model.contact_geoms:
contact_set.add(g2)
elif g2 in model.contact_geoms and g1 not in model.contact_geoms:
contact_set.add(g1)
return contact_set
def add_observable(self, observable):
"""
Adds an observable to this environment.
Args:
observable (Observable): Observable instance.
"""
assert observable.name not in self._observables,\
"Observable name {} is already associated with an existing observable! Use modify_observable(...) " \
"to modify a pre-existing observable.".format(observable.name)
self._observables[observable.name] = observable
def modify_observable(self, observable_name, attribute, modifier):
"""
Modifies observable with associated name @observable_name, replacing the given @attribute with @modifier.
Args:
observable_name (str): Observable to modify
attribute (str): Observable attribute to modify.
Options are {`'sensor'`, `'corrupter'`,`'filter'`, `'delayer'`, `'sampling_rate'`,
`'enabled'`, `'active'`}
modifier (any): New function / value to replace with for observable. If a function, new signature should
match the function being replaced.
"""
# Find the observable
assert observable_name in self._observables, "No valid observable with name {} found. Options are: {}".\
format(observable_name, self.observation_names)
obs = self._observables[observable_name]
# replace attribute accordingly
if attribute == "sensor":
obs.set_sensor(modifier)
elif attribute == "corrupter":
obs.set_corrupter(modifier)
elif attribute == "filter":
obs.set_filter(modifier)
elif attribute == "delayer":
obs.set_delayer(modifier)
elif attribute == "sampling_rate":
obs.set_sampling_rate(modifier)
elif attribute == "enabled":
obs.set_enabled(modifier)
elif attribute == "active":
obs.set_active(modifier)
else:
# Invalid attribute specified
raise ValueError("Invalid observable attribute specified. Requested: {}, valid options are {}".
format(attribute, {"sensor", "corrupter", "filter", "delayer",
"sampling_rate", "enabled", "active"}))
def _check_success(self):
"""
Checks if the task has been completed. Should be implemented by subclasses
Returns:
bool: True if the task has been completed
"""
raise NotImplementedError
def _destroy_viewer(self):
"""
Destroys the current mujoco renderer instance if it exists
"""
# 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()
@property
def observation_modalities(self):
"""
Modalities for this environment's observations
Returns:
set: All observation modalities
"""
return set([observable.modality for observable in self._observables.values()])
@property
def observation_names(self):
"""
Grabs all names for this environment's observables
Returns:
set: All observation names
"""
return set(self._observables.keys())
@property
def enabled_observables(self):
"""
Grabs all names of enabled observables for this environment. An observable is considered enabled if its values
are being continually computed / updated at each simulation timestep.
Returns:
set: All enabled observation names
"""
return set([name for name, observable in self._observables.items() if observable.is_enabled()])
@property
def active_observables(self):
"""
Grabs all names of active observables for this environment. An observable is considered active if its value is
being returned in the observation dict from _get_observations() call or from the step() call (assuming this
observable is enabled).
Returns:
set: All active observation names
"""
return set([name for name, observable in self._observables.items() if observable.is_active()])
@property
def _visualizations(self):
"""
Visualization keywords for this environment
Returns:
set: All components that can be individually visualized for this environment
"""
return {"env"}
@property
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
@property
def action_dim(self):
"""
Size of the action space
Returns:
int: Action space dimension
"""
raise NotImplementedError
class WindySlope(gym.Env):
def __init__(self,
model,
mode,
hertz=25,
should_render=True,
should_screenshot=False):
self.hertz = hertz
self.steps = 0
self.should_render = should_render
self.should_screenshot = should_screenshot
self.nsubsteps = int(MAX_TIME / model.opt.timestep / 100)
self.viewer = None
self.model = model
self.mode = mode
self.enabled = True
self.metadata = {'render.modes': 'rgb_array'}
self.should_record = False
def close(self):
pass
@property
def observation_space(self):
return Box(low=-np.inf, high=np.inf, shape=(18, ))
@property
def action_space(self):
return Box(low=-np.inf, high=np.inf, shape=(0, ))
@property
def model(self):
return self._model
@model.setter
def model(self, model):
self._model = model
self.sim = MjSim(model)
self.data = self.sim.data
if self.should_render:
if self.viewer:
self.viewer.sim = sim
return
self.viewer = MjViewer(self.sim)
self.viewer.cam.azimuth = 45
self.viewer.cam.elevation = -20
self.viewer.cam.distance = 25
self.viewer.cam.lookat[:] = [0, 0, -2]
self.viewer.scn.flags[3] = 0
def reset(self):
self.sim.reset()
self.steps = 0
self.sim.forward()
obs = self.get_observations(self.model, self.data)
return obs
def get_observations(self, model, data):
self.sim.forward()
obs = []
name = 'box'
pos = data.get_body_xpos(name)
xmat = data.get_body_xmat(name).reshape(-1)
velp = data.get_body_xvelp(name)
velr = data.get_body_xvelr(name)
for x in [pos, xmat, velp, velr]:
obs.extend(x.copy())
obs = np.array(obs, dtype=np.float32)
return obs
def screenshot(self, image_path):
self.viewer.hide_overlay = True
self.viewer.render()
width, height = 2560, 1440
#width, height = 1,1
img = self.viewer.read_pixels(width, height, depth=False)
# original image is upside-down, so flip it
img = img[::-1, :, :]
imageio.imwrite(image_path, img)
def step(self, action):
nsubsteps = self.nsubsteps
for _ in range(nsubsteps):
self.sim.step()
self.render()
self.steps += 1
return self.get_observations(self.model,
self.data), 1, self.steps == 100, {}
def render(self, mode=None):
if self.should_render:
self.viewer._overlay.clear()
def nothing():
return
self.viewer._create_full_overlay = nothing
wind = model.opt.wind[0]
self.viewer.add_overlay(1, "Wind", "{:.2f}".format(wind))
self.viewer.render()
if self.should_record:
width, height = 2560, 1440
img = self.viewer.read_pixels(width, height, depth=False)
# original image is upside-down, so flip it
img = img[::-1, :, :]
return img
def euler2quat(self, euler):
euler = np.asarray(euler, dtype=np.float64)
assert euler.shape[-1] == 3, "Invalid shape euler {}".format(euler)
ai, aj, ak = euler[..., 2] / 2, -euler[..., 1] / 2, euler[..., 0] / 2
si, sj, sk = np.sin(ai), np.sin(aj), np.sin(ak)
ci, cj, ck = np.cos(ai), np.cos(aj), np.cos(ak)
cc, cs = ci * ck, ci * sk
sc, ss = si * ck, si * sk
quat = np.empty(euler.shape[:-1] + (4, ), dtype=np.float64)
quat[..., 0] = cj * cc + sj * ss
quat[..., 3] = cj * sc - sj * cs
quat[..., 2] = -(cj * ss + sj * cc)
quat[..., 1] = cj * cs - sj * sc
return quat
def degrees2radians(self, degrees):
return degrees * np.pi / 180
model = load_model_from_xml(MODEL_XML)
sim = MjSim(model)
viewer = MjViewer(sim)
states = [{'box:x': +0.8, 'box:y': +0.8},
{'box:x': -0.8, 'box:y': +0.8},
{'box:x': -0.8, 'box:y': -0.8},
{'box:x': +0.8, 'box:y': -0.8},
{'box:x': +0.0, 'box:y': +0.0}]
# MjModel.joint_name2id returns the index of a joint in
# MjData.qpos.
x_joint_i = sim.model.get_joint_qpos_addr("box:x")
y_joint_i = sim.model.get_joint_qpos_addr("box:y")
print_box_xpos(sim)
while True:
for state in states:
sim_state = sim.get_state()
sim_state.qpos[x_joint_i] = state["box:x"]
sim_state.qpos[y_joint_i] = state["box:y"]
sim.set_state(sim_state)
sim.forward()
print("updated state to", state)
print_box_xpos(sim)
viewer.render()
if os.getenv('TESTING') is not None:
break
class MujocoEnv(gym.Env):
"""Superclass for all MuJoCo environments.
"""
def __init__(self, model_path, frame_skip):
if model_path.startswith("/"):
fullpath = model_path
else:
fullpath = os.path.join(os.path.dirname(__file__), "assets", model_path)
if not path.exists(fullpath):
raise IOError("File %s does not exist" % fullpath)
self.frame_skip = frame_skip
self.model = load_model_from_path(fullpath)
self.sim = MjSim(self.model)
self.data = self.sim.data
self.metadata = {
'render.modes': ['human', 'rgb_array'],
'video.frames_per_second': int(np.round(1.0 / self.dt))
}
self.mujoco_render_frames = False
self.init_qpos = self.data.qpos.ravel().copy()
self.init_qvel = self.data.qvel.ravel().copy()
observation, _reward, done, _info = self._step(np.zeros(self.model.nu))
assert not done
self.obs_dim = np.sum([o.size for o in observation]) if type(observation) is tuple else observation.size
bounds = self.model.actuator_ctrlrange.copy()
low = bounds[:, 0]
high = bounds[:, 1]
self.action_space = spaces.Box(low, high)
high = np.inf*np.ones(self.obs_dim)
low = -high
self.observation_space = spaces.Box(low, high)
self._seed()
def _seed(self, seed=None):
self.np_random, seed = seeding.np_random(seed)
return [seed]
# methods to override:
# ----------------------------
def reset_model(self):
"""
Reset the robot degrees of freedom (qpos and qvel).
Implement this in each subclass.
"""
raise NotImplementedError
def mj_viewer_setup(self):
"""
Due to specifics of new mujoco rendering, the standard viewer cannot be used
with this set-up. Instead we use this mujoco specific function.
"""
pass
def viewer_setup(self):
"""
Does not work. Use mj_viewer_setup() instead
"""
pass
# -----------------------------
def _reset(self):
self.sim.reset()
self.sim.forward()
ob = self.reset_model()
return ob
def set_state(self, qpos, qvel):
assert qpos.shape == (self.model.nq,) and qvel.shape == (self.model.nv,)
state = self.sim.get_state()
for i in range(self.model.nq):
state.qpos[i] = qpos[i]
for i in range(self.model.nv):
state.qvel[i] = qvel[i]
self.sim.set_state(state)
self.sim.forward()
@property
def dt(self):
return self.model.opt.timestep * self.frame_skip
def do_simulation(self, ctrl, n_frames):
for i in range(self.model.nu):
self.sim.data.ctrl[i] = ctrl[i]
for _ in range(n_frames):
self.sim.step()
if self.mujoco_render_frames is True:
self.mj_render()
def mj_render(self):
try:
self.viewer.render()
except:
self.mj_viewer_setup()
self.viewer._run_speed = 1.0
#self.viewer._run_speed /= self.frame_skip
self.viewer.render()
def _get_viewer(self):
return None
def state_vector(self):
state = self.sim.get_state()
return np.concatenate([
state.qpos.flat, state.qvel.flat])
# -----------------------------
def visualize_policy(self, policy, horizon=1000, num_episodes=1, mode='exploration'):
self.mujoco_render_frames = True
for ep in range(num_episodes):
o = self.reset()
d = False
t = 0
while t < horizon and d is False:
a = policy.get_action(o)[0] if mode == 'exploration' else policy.get_action(o)[1]['evaluation']
o, r, d, _ = self.step(a)
t = t+1
self.mujoco_render_frames = False
def visualize_policy_offscreen(self, policy, horizon=1000,
num_episodes=1,
mode='exploration',
save_loc='/tmp/',
filename='newvid',
camera_name=None):
import skvideo.io
for ep in range(num_episodes):
print("Episode %d: rendering offline " % ep, end='', flush=True)
o = self.reset()
d = False
t = 0
arrs = []
t0 = timer.time()
while t < horizon and d is False:
a = policy.get_action(o)[0] if mode == 'exploration' else policy.get_action(o)[1]['evaluation']
o, r, d, _ = self.step(a)
t = t+1
curr_frame = self.sim.render(width=640, height=480, mode='offscreen',
camera_name=camera_name, device_id=0)
arrs.append(curr_frame[::-1,:,:])
print(t, end=', ', flush=True)
file_name = save_loc + filename + str(ep) + ".mp4"
skvideo.io.vwrite( file_name, np.asarray(arrs))
print("saved", file_name)
t1 = timer.time()
print("time taken = %f"% (t1-t0))
def test_arrays_of_objs():
model = load_model_from_xml(BASIC_MODEL_XML)
sim = MjSim(model)
sim.forward()
renderer = cymj.MjRenderContext(sim, offscreen=True)
assert len(renderer.scn.camera) == 2, "Expecting scn.camera to be available"
def test_jacobians():
xml = """
<mujoco>
<worldbody>
<body name="body1" pos="0 0 0">
<joint axis="1 0 0" name="a" pos="0 0 0" type="hinge"/>
<geom name="geom1" pos="0 0 0" size="1.0"/>
<body name="body2" pos="0 0 1">
<joint name="b" axis="1 0 0" pos="0 0 1" type="hinge"/>
<geom name="geom2" pos="1 1 1" size="0.5"/>
<site name="target" size="0.1"/>
</body>
</body>
</worldbody>
<actuator>
<motor joint="a"/>
<motor joint="b"/>
</actuator>
</mujoco>
"""
model = load_model_from_xml(xml)
sim = MjSim(model)
sim.reset()
# After reset jacobians are all zeros
target_jacp = np.zeros(3 * sim.model.nv)
sim.data.get_site_jacp('target', jacp=target_jacp)
np.testing.assert_allclose(target_jacp, np.zeros(3 * sim.model.nv))
# After first forward, jacobians are real
sim.forward()
sim.data.get_site_jacp('target', jacp=target_jacp)
target_test = np.array([0, 0, -1, 1, 0, 0])
np.testing.assert_allclose(target_jacp, target_test)
# Should be unchanged after steps (zero action)
for _ in range(2):
sim.step()
sim.forward()
sim.data.get_site_jacp('target', jacp=target_jacp)
assert np.linalg.norm(target_jacp - target_test) < 1e-3
# Apply a very large action, ensure jacobian unchanged after step
sim.reset()
sim.forward()
sim.data.ctrl[:] = np.ones(sim.model.nu) * 1e9
sim.step()
sim.data.get_site_jacp('target', jacp=target_jacp)
np.testing.assert_allclose(target_jacp, target_test)
# After large action, ensure jacobian changed after forward
sim.forward()
sim.data.get_site_jacp('target', jacp=target_jacp)
assert not np.allclose(target_jacp, target_test)
# Test the `site_jacp` property, which gets all at once
np.testing.assert_allclose(target_jacp, sim.data.site_jacp[0])
# Test not passing in array
sim.reset()
sim.forward()
target_jacp = sim.data.get_site_jacp('target')
np.testing.assert_allclose(target_jacp, target_test)
# Test passing in bad array (long instead of double)
target_jacp = np.zeros(3 * sim.model.nv, dtype=np.long)
with pytest.raises(ValueError):
sim.data.get_site_jacp('target', jacp=target_jacp)
# Test rotation jacobian - like above but 'jacr' instead of 'jacp'
# After reset jacobians are all zeros
sim.reset()
target_jacr = np.zeros(3 * sim.model.nv)
sim.data.get_site_jacr('target', jacr=target_jacr)
np.testing.assert_allclose(target_jacr, np.zeros(3 * sim.model.nv))
# After first forward, jacobians are real
sim.forward()
sim.data.get_site_jacr('target', jacr=target_jacr)
target_test = np.array([1, 1, 0, 0, 0, 0])
# Test allocating dedicated array
target_jacr = sim.data.get_site_jacr('target')
np.testing.assert_allclose(target_jacr, target_test)
# Test the batch getter (all sites at once)
np.testing.assert_allclose(target_jacr, sim.data.site_jacr[0])
# Test passing in bad array
target_jacr = np.zeros(3 * sim.model.nv, dtype=np.long)
with pytest.raises(ValueError):
sim.data.get_site_jacr('target', jacr=target_jacr)
model = load_model_from_xml(MODEL_XML)
sim = MjSim(model)
viewer = MjViewer(sim)
states = [{'box:x': +0.8, 'box:y': +0.8},
{'box:x': -0.8, 'box:y': +0.8},
{'box:x': -0.8, 'box:y': -0.8},
{'box:x': +0.8, 'box:y': -0.8},
{'box:x': +0.0, 'box:y': +0.0}]
# MjModel.joint_name2id returns the index of a joint in
# MjData.qpos.
x_joint_i = sim.model.get_joint_qpos_addr("box:x")
y_joint_i = sim.model.get_joint_qpos_addr("box:y")
print(sim.model.get_joint_names())
asdf
print_box_xpos(sim)
while True:
for state in states:
sim_state = sim.get_state()
sim_state.qpos[x_joint_i] = state["box:x"]
sim_state.qpos[y_joint_i] = state["box:y"]
sim.set_state(sim_state)
sim.forward()
print("updated state to", state)
print_box_xpos(sim)
viewer.render()
if os.getenv('TESTING') is not None:
break
class lab_env():
def __init__(self, env, args):
#super(lab_env, self).__init__(env)
# The real-world simulator
self.model = load_model_from_path('xmls/lab_env.xml')
self.sim = MjSim(self.model)
# Used to locate the gripper
self.model2 = load_model_from_path('xmls/light_env.xml')
self.sim2 = MjSim(self.model2)
def reset(self, task_id):
self.task = task_id
self.grasping = -1
self.last_grasp = -1
# Configure gravity
for i in range(4):
self.sim.data.ctrl[i] = -1
# Configure joint positions
for i in range(42):
self.sim.data.qpos[i] = initial_pos[i]
for i in range(3):
self.sim.data.qpos[i] = joint_pos[task_id][i]
self.pos_forward()
self.sim.forward()
remapped_pos = [
remap(self.sim.data.qpos[0], -30 / 180 * math.pi,
45 / 180 * math.pi, -1, 1),
remap(self.sim.data.qpos[1], -105 / 180 * math.pi,
-50 / 180 * math.pi, -1, 1),
remap(self.sim.data.qpos[2], 0 / 180 * math.pi,
180 / 180 * math.pi, -1, 1), 0
]
return (remapped_pos, ) + self.get_state()
def step(self, action):
self.sim.data.qpos[0] = remap(action[0], -1, 1, -30 / 180 * math.pi,
45 / 180 * math.pi)
self.sim.data.qpos[1] = remap(action[1], -1, 1, -105 / 180 * math.pi,
-50 / 180 * math.pi)
self.sim.data.qpos[2] = remap(action[2], -1, 1, 0 / 180 * math.pi,
180 / 180 * math.pi)
self.pos_forward()
self.sim.forward()
if action[3] < self.last_grasp or self.grasping == -1:
t = int(remap(action[3], -1, 1, 0, grasp_steps))
for i in range(6, 14):
self.sim.data.qpos[i] = 0
self.sim.forward()
self.grasping = -1
self.sim.data.ctrl[4] = 1
self.sim.data.ctrl[5] = 1
backup = [
self.sim.data.qpos[i]
for i in [15, 16, 22, 23, 29, 30, 36, 37]
]
for i in range(t):
self.sim.step()
stop = False
for j in range(4):
if self.sim.data.sensordata[j] > sensor_threshold:
self.grasping = j
self.pickuppos = [
self.sim.data.qpos[i]
for i in (list(range(6)) + list(
range(14 + 7 * self.grasping, 21 +
7 * self.grasping)))
]
stop = True
break
for i in range(4):
for j in range(2):
self.sim.data.qpos[15 + 7 * i + j] = backup[i * 2 + j]
if stop:
break
self.gripper_sync()
self.sim.forward()
self.sim.data.ctrl[4] = 0
self.sim.data.ctrl[5] = 0
self.last_grasp = action[3]
return self.get_state()
# Ensure that the gripper is always heading down and is parallar to the desk edge
def pos_forward(self):
self.sim.data.qpos[
3] = math.pi * 1.5 - self.sim.data.qpos[1] - self.sim.data.qpos[2]
self.sim.data.qpos[4] = math.pi * 1.5
self.sim.data.qpos[5] = math.pi * 1.25 + self.sim.data.qpos[0]
self.gripper_sync()
if self.grasping != -1:
current_xyz = pos_to_xyz(self.sim.data.qpos[0:6])
old_xyz = pos_to_xyz(self.pickuppos[0:6])
for i in range(3):
self.sim.data.qpos[
14 + 7 * self.grasping +
i] = self.pickuppos[6 + i] + current_xyz[i] - old_xyz[i]
'''
old_quat = quaternion(self.pickuppos[9], vector(self.pickuppos[10], self.pickuppos[11], self.pickuppos[12]))
rotate_quat = quaternion(math.cos(self.sim.data.qpos[0] - self.pickuppos[0]), vector(0, 0, math.sin(self.sim.data.qpos[0] - self.pickuppos[0])))
new_quat = rotate_quat.mul(old_quat)
self.sim.data.qpos[17 + 7 * self.grasping] = new_quat.w
self.sim.data.qpos[18 + 7 * self.grasping] = new_quat.v.x
self.sim.data.qpos[19 + 7 * self.grasping] = new_quat.v.y
self.sim.data.qpos[20 + 7 * self.grasping] = new_quat.v.z
'''
def gripper_sync(self):
self.sim.data.qpos[9] = gripper_consistent(self.sim.data.qpos[6:9])
self.sim.data.qpos[13] = gripper_consistent(self.sim.data.qpos[10:13])
def get_state(self):
sync(self.sim, self.sim2, 6)
# Locate the gripper, render twice to overcome bugs in mujoco
image_1 = copy.deepcopy(
self.sim.render(width=960,
height=720,
camera_name='workbench_camera'))
image_1 = copy.deepcopy(
self.sim.render(width=960,
height=720,
camera_name='workbench_camera'))
image_2 = copy.deepcopy(
self.sim.render(width=960, height=720, camera_name='upper_camera'))
image_2 = copy.deepcopy(
self.sim.render(width=960, height=720, camera_name='upper_camera'))
image_3 = copy.deepcopy(
self.sim2.render(width=960,
height=720,
camera_name='workbench_camera'))
image_3 = copy.deepcopy(
self.sim2.render(width=960,
height=720,
camera_name='workbench_camera'))
x1, y1 = get_x_y(image_3)
image_4 = copy.deepcopy(
self.sim2.render(width=960, height=720,
camera_name='upper_camera'))
image_4 = copy.deepcopy(
self.sim2.render(width=960, height=720,
camera_name='upper_camera'))
x2, y2 = get_x_y(image_4)
# Crop gripper images and add noise
image_1 = cv2.GaussianBlur(
gaussian_noise(
crop(image_1, x1 + fig_size_1[0] // 2, y1, *fig_size_1),
*gaussian_noise_parameters),
*gaussian_blur_prarmeters).transpose((2, 0, 1))
image_2 = cv2.GaussianBlur(
gaussian_noise(
crop(image_2, x2 + fig_size_2[0] // 2, y2, *fig_size_2),
*gaussian_noise_parameters),
*gaussian_blur_prarmeters).transpose((2, 0, 1))
danger = int(self.safety_check() or math.isnan(x1) or math.isnan(y1)
or math.isnan(x2) or math.isnan(y2))
return [x1, y1, x2, y2,
int(self.grasping == self.task), danger], (image_1, image_2)
def safety_check(self):
# return 0 if safe, otherwise 1
backup = [self.sim.data.qpos[i] for i in range(14)]
self.sim.step()
s = 0
for i in range(6):
s += abs(backup[i] - self.sim.data.qpos[i])
self.sim.data.qpos[i] = backup[i]
return s > safety_threshold
def force_close(self):
for i in range(2):
for j in range(3):
self.sim.data.qpos[6 + i * 4 + j] = closed_pos[j]
self.gripper_sync()
self.sim.forward()
class MujocoEnv(gym.Env):
"""Superclass for all MuJoCo environments.
"""
def __init__(self,
env_name,
rand_seed,
maximum_length,
model_path,
frame_skip,
misc_info={}):
self._env_name = env_name
self._rand_seed = rand_seed
self._maximum_length = maximum_length
self._misc_info = misc_info
if model_path.startswith("/"):
fullpath = model_path
else:
fullpath = os.path.join(os.path.dirname(__file__), "assets",
model_path)
if not path.exists(fullpath):
raise IOError("File %s does not exist" % fullpath)
self.frame_skip = frame_skip
self.model = load_model_from_path(fullpath)
self.sim = MjSim(self.model)
self.data = self.sim.data
self.metadata = {
'render.modes': ['human', 'rgb_array'],
'video.frames_per_second': int(np.round(1.0 / self.dt))
}
self.mujoco_render_frames = False
self.init_qpos = self.data.qpos.ravel().copy()
self.init_qvel = self.data.qvel.ravel().copy()
# Why is this here? Causes errors
#observation, _reward, done, _info = self._step(np.zeros(self.model.nu))
#assert not done
#self.obs_dim = np.sum([o.size for o in observation]) if type(observation) is tuple else observation.size
bounds = self.model.actuator_ctrlrange.copy()
low = bounds[:, 0]
high = bounds[:, 1]
self.action_space = spaces.Box(low, high)
# annoying should replace
high = np.inf * np.ones(self.obs_dim)
low = -high
self.observation_space = spaces.Box(low, high)
self._seed(self._rand_seed)
def _seed(self, seed=None):
self.np_random, seed = seeding.np_random(self._rand_seed)
return [seed]
# methods to override:
# ----------------------------
def reset_model(self):
"""
Reset the robot degrees of freedom (qpos and qvel).
Implement this in each subclass.
"""
raise NotImplementedError
def mj_viewer_setup(self):
"""
Due to specifics of new mujoco rendering, the standard viewer cannot be used
with this set-up. Instead we use this mujoco specific function.
"""
pass
def get_env_state(self):
"""
Get full state of the environment beyond qpos and qvel
For example, if targets are defined using sites, this function should also
contain location of the sites (which are not included in qpos).
Must return a dictionary that can be used in the set_env_state function
"""
raise NotImplementedError
def set_env_state(self, state):
"""
Uses the state dictionary to set the state of the world
"""
raise NotImplementedError
# -----------------------------
def reset(self):
self.sim.reset()
self.sim.forward()
ob = self.reset_model()
return ob, 0, False, {}
def reset_soft(self, seed=None):
return self._old_obs, 0, False, {}
def set_state(self, qpos, qvel):
assert qpos.shape == (self.model.nq, ) and qvel.shape == (
self.model.nv, )
state = self.sim.get_state()
for i in range(self.model.nq):
state.qpos[i] = qpos[i]
for i in range(self.model.nv):
state.qvel[i] = qvel[i]
self.sim.set_state(state)
self.sim.forward()
@property
def dt(self):
return self.model.opt.timestep * self.frame_skip
def do_simulation(self, ctrl, n_frames):
for i in range(self.model.nu):
self.sim.data.ctrl[i] = ctrl[i]
for _ in range(n_frames):
self.sim.step()
if self.mujoco_render_frames is True:
self.mj_render()
def mj_render(self):
try:
self.viewer.render()
except:
self.mj_viewer_setup()
self.viewer._run_speed = 0.5
#self.viewer._run_speed /= self.frame_skip
self.viewer.render()
def _get_viewer(self):
return None
def state_vector(self):
state = self.sim.get_state()
return np.concatenate([state.qpos.flat, state.qvel.flat])
# -----------------------------
def visualize_policy(self,
policy,
horizon=1000,
num_episodes=1,
mode='exploration'):
self.mujoco_render_frames = True
for ep in range(num_episodes):
o = self.reset()
d = False
t = 0
while t < horizon and d is False:
a = policy.get_action(
o)[0] if mode == 'exploration' else policy.get_action(
o)[1]['evaluation']
o, r, d, _ = self.step(a)
t = t + 1
self.mujoco_render_frames = False
def visualize_policy_offscreen(self,
policy,
horizon=1000,
num_episodes=1,
frame_size=(640, 480),
mode='exploration',
save_loc='./tmp/',
filename='newvid',
it=0,
camera_name=None):
for ep in range(num_episodes):
print("Episode %d: rendering offline " % ep, end='', flush=True)
o, *_ = self.reset()
d = False
t = 0
arrs = []
t0 = timer.time()
while t < horizon and d is False:
a = policy(o)
o, r, d, _ = self.step(a)
t = t + 1
curr_frame = self.sim.render(width=frame_size[0],
height=frame_size[1],
mode='offscreen',
camera_name=camera_name,
device_id=0)
arrs.append(curr_frame[::-1, :, :])
print(t, end=', ', flush=True)
file_name = save_loc + filename + str(ep) + str(it) + ".mp4"
if not os.path.exists(save_loc):
os.makedirs(save_loc)
imageio.mimwrite(file_name, np.asarray(arrs), fps=10.0)
print("saved", file_name)
t1 = timer.time()
print("time taken = %f" % (t1 - t0))