def test_userdata(self):
xml = '''
<mujoco>
<size nuserdata="{}"/>
<worldbody>
<body pos="0 0 0">
<joint type="free"/>
<geom type="sphere" size=".1"/>
</body>
</worldbody>
</mujoco>
'''
model = load_model_from_xml(xml.format(1))
assert model.nuserdata == 1, 'bad nuserdata {}'.format(model.nuserdata)
model = load_model_from_xml(xml.format(10))
assert model.nuserdata == 10, 'bad nuserdata {}'.format(model.nuserdata)
sim = MjSim(model)
data = sim.data
assert data.userdata[0] == 0, 'bad userdata {}'.format(data.userdata)
data.userdata[0] = 1
assert data.userdata[0] == 1, 'bad userdata {}'.format(data.userdata)
# Check that we throw an assert if there's not enough userdata
model = load_model_from_xml(xml.format(0))
with self.assertRaises(AssertionError):
model.set_userdata_names(['foo'])
# Doesn't throw assert
model = load_model_from_xml(xml.format(1))
model.set_userdata_names(['foo'])
with self.assertRaises(AssertionError):
model.set_userdata_names(['foo', 'bar'])
def test_model_save_load():
model = load_model_from_xml(BASIC_MODEL_XML)
xml_from_model = model.get_xml()
model_from_xml = load_model_from_xml(xml_from_model)
assert(xml_from_model == model_from_xml.get_xml())
mjb_from_model = model.get_mjb()
model_from_mjb = load_model_from_mjb(mjb_from_model)
assert(mjb_from_model == model_from_mjb.get_mjb())
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_concurrent_rendering():
'''Best-effort testing that concurrent multi-threaded rendering works.
The test has no guarantees around being deterministic, but if it fails
you know something is wrong with concurrent rendering. If it passes,
things are probably working.'''
err = None
def func(sim, event):
event.wait()
sim.data.qpos[:] = 0.0
sim.forward()
img1 = sim.render(width=40, height=40, camera_name="camera1")
img2 = sim.render(width=40, height=40, camera_name="camera2")
try:
assert np.sum(img1[:]) == 23255
assert np.sum(img2[:]) == 12007
except Exception as e:
nonlocal err
err = e
model = load_model_from_xml(BASIC_MODEL_XML)
sim = MjSim(model)
sim.render(100, 100)
event = Event()
threads = []
for _ in range(100):
thread = Thread(target=func, args=(sim, event))
threads.append(thread)
thread.start()
event.set()
for thread in threads:
thread.join()
assert err is None, "Exception: %s" % (str(err))
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_ignore_mujoco_warnings():
# Two boxes on a plane need more than 1 contact (nconmax)
xml = '''
<mujoco>
<size nconmax="1"/>
<worldbody>
<geom type="plane" size="1 1 0.1"/>
<body pos="1 0 1"> <joint type="free"/> <geom size="1"/> </body>
<body pos="0 1 1"> <joint type="free"/> <geom size="1"/> </body>
</worldbody>
</mujoco>
'''
model = load_model_from_xml(xml)
sim = MjSim(model)
sim.reset()
with ignore_mujoco_warnings():
# This should raise an exception due to the mujoco warning callback,
# but it's suppressed by the context manager.
sim.step()
sim.reset()
with pytest.raises(Exception):
# test to make sure previous warning callback restored.
sim.step()
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_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_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_mj_proptional_only(): model = load_model_from_xml(MODEL_XML.format(actuator=P_ONLY_ACTUATOR)) sim = MjSim(model) cymj.set_pid_control(sim.model, sim.data) model2 = load_model_from_xml(MODEL_XML.format(actuator=POSITION_ACTUATOR)) sim2 = MjSim(model2) init_pos = 0.1 * (random.random() - 0.5) sim.data.qpos[0] = sim2.data.qpos[0] = init_pos sim.data.ctrl[0] = sim2.data.ctrl[0] = 0 for i in range(20): print(i, sim.data.qpos[0], sim2.data.qpos[0]) sim.step() sim2.step() assert abs(sim.data.qpos[0] - sim2.data.qpos[0]) <= 1e-5, "%d step violates" % i
def load_model(self, path):
self.tree = ET.parse(path)
root = self.tree.getroot()
root = self.init_task(root)
with io.StringIO() as string:
string.write(ET.tostring(root, encoding="unicode"))
model = mujoco_py.load_model_from_xml(string.getvalue())
return model
def test_mj_sim_pool_basics():
model = load_model_from_xml(BASIC_MODEL_XML)
sims = [MjSim(model) for _ in range(2)]
sim_pool = MjSimPool(sims, nsubsteps=2)
sim_pool.reset()
sim_pool.step()
sim_pool.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_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_mj_sim_pool_basics():
model = load_model_from_xml(BASIC_MODEL_XML)
sims = [MjSim(model) for _ in range(2)]
sim_pool = MjSimPool(sims, nsubsteps=2)
sim_pool.reset()
sim_pool.step()
sim_pool.forward()
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_cascaded_pdpi():
"""
This tests using a PD-PI loop with the cascaded controller. This is achieved by setting the
integral time constant and clamp equal to zero.
Here we setup two sims (CASCADED_PDPI_ACTUATOR and CASCADED_PDPI_ACTUATOR_NO_D), one with
derivative gain and one without. With the exception of the derivative term, all other gain
parameters are the same. The goal is to show the stability added by the derivative term.
"""
random.seed(30)
xml = MODEL_XML.format(actuator=CASCADED_PDPI_ACTUATOR, nuser_actuator=1)
model = load_model_from_xml(xml)
sim = MjSim(model)
cymj.set_pid_control(sim.model, sim.data)
"""
sim2 uses the same Kp gain on the position loop as sim but does not have any derivative gain.
It is expected that given this Kp gain and no derivative gain, the system will be unstable
and fail to hold the desired position.
"""
xml = MODEL_XML.format(actuator=CASCADED_PDPI_ACTUATOR_NO_D, nuser_actuator=1)
model2 = load_model_from_xml(xml)
sim2 = MjSim(model2)
cymj.set_pid_control(sim2.model, sim2.data)
# Perturbation of pole to be unbalanced
init_pos = 0.1 * (random.random() - 1.0)
print('init pos', init_pos)
sim.data.qpos[0] = sim2.data.qpos[0] = init_pos
desired_pos = 0.0
sim.data.ctrl[0] = sim2.data.ctrl[0] = desired_pos
print('desired position:', desired_pos)
max_torque = 0
for _ in range(2000):
sim.step()
sim2.step()
if abs(sim.data.actuator_force[0]) > max_torque:
max_torque = abs(sim.data.actuator_force[0])
print('final pos', sim.data.qpos[0])
assert abs(sim2.data.qpos[0] - desired_pos) > 1e-3 # Verify instability without D term
assert abs(sim.data.qpos[0] - desired_pos) < 1e-3 # Verify stability with D term
assert max_torque <= 3 # Torque limit set on the actuator
def test_many_sims_rendering():
model = load_model_from_xml(BASIC_MODEL_XML)
sims = [MjSim(model) for _ in range(5)]
pool = MjSimPool(sims)
pool.forward()
for sim in sims:
img, depth = sim.render(200, 200, depth=True)
assert img.shape == (200, 200, 3)
compare_imgs(img, 'test_rendering.freecam.png')
def experiment(variant):
MODEL_XML = """
<?xml version="1.0" ?>
<mujoco>
<option timestep="0.005" />
<worldbody>
<body name="robot" pos="0 0 1.2">
<joint axis="1 0 0" damping="0.1" name="slide0" pos="0 0 0" type="slide"/>
<joint axis="0 1 0" damping="0.1" name="slide1" pos="0 0 0" type="slide"/>
<joint axis="0 0 1" damping="1" name="slide2" pos="0 0 0" type="slide"/>
<geom mass="1.0" pos="0 0 0" rgba="1 0 0 1" size="0.15" type="sphere"/>
<camera euler="0 0 0" fovy="40" name="rgb" pos="0 0 2.5"></camera>
</body>
<body mocap="true" name="mocap" pos="0.5 0.5 0.5">
<geom conaffinity="0" contype="0" pos="0 0 0" rgba="1.0 1.0 1.0 0.5" size="0.1 0.1 0.1" type="box"></geom>
<geom conaffinity="0" contype="0" pos="0 0 0" rgba="1.0 1.0 1.0 0.5" size="0.2 0.2 0.05" type="box"></geom>
</body>
<body name="cylinder" pos="0.1 0.1 0.2">
<geom mass="1" size="0.15 0.15" type="cylinder"/>
<joint axis="1 0 0" name="cylinder:slidex" type="slide"/>
<joint axis="0 1 0" name="cylinder:slidey" type="slide"/>
</body>
<body name="box" pos="-0.8 0 0.2">
<geom mass="0.1" size="0.15 0.15 0.15" type="box"/>
</body>
<body name="floor" pos="0 0 0.025">
<geom condim="3" size="1.0 1.0 0.02" rgba="0 1 0 1" type="box"/>
</body>
</worldbody>
<actuator>
<motor gear="2000.0" joint="slide0"/>
<motor gear="2000.0" joint="slide1"/>
</actuator>
</mujoco>
"""
model = load_model_from_xml(MODEL_XML)
sim = MjSim(model)
# viewer = MjViewer(sim)
# t = 0
# while True:
# sim.data.ctrl[0] = math.cos(t / 10.) * 0.01
# sim.data.ctrl[1] = math.sin(t / 10.) * 0.01
# t += 1
# sim.step()
# viewer.render()
# if t > 100 and os.getenv('TESTING') is not None:
# break
renderer = MjBatchRenderer(100, 100, use_cuda=True)
renderer.render(sim)
renderer.map()
image = renderer.read()
renderer.unmap()
print(image)
print(image[0].mean())
def test_many_sims_rendering():
model = load_model_from_xml(BASIC_MODEL_XML)
sims = [MjSim(model) for _ in range(5)]
pool = MjSimPool(sims)
pool.forward()
for sim in sims:
img, depth = sim.render(200, 200, depth=True)
assert img.shape == (200, 200, 3)
compare_imgs(img, 'test_rendering.freecam.png')
def mp_test_create_destroy():
model = load_model_from_xml(BASIC_MODEL_XML)
pool = MjRenderPool(model, n_workers=2)
del pool
# closing before deleting should be fine too
pool = MjRenderPool(model, n_workers=2)
pool.close()
del pool
def create_model(xml):
model = load_model_from_xml(postprocess_model_xml(xml))
model.vis.quality.offsamples = 8
sim = MjSim(model)
render_context = MjRenderContextOffscreen(sim)
render_context.vopt.geomgroup[0] = 0
render_context.vopt.geomgroup[1] = 1
sim.add_render_context(render_context)
return sim
def test_sim_save():
model = load_model_from_xml(BASIC_MODEL_XML)
assert model.nkey == 0
sim = MjSim(model)
with StringIO() as f:
sim.save(f)
f.seek(0)
loaded_model = load_model_from_xml(f.read())
assert loaded_model.nkey == 1
with BytesIO() as f:
sim.save(f, format='mjb')
f.seek(0)
loaded_model = load_model_from_mjb(f.read())
assert loaded_model.nkey == 1
def test_hello(self):
fn = '''
#include <stdio.h>
void fun(const mjModel* m, mjData* d) {
printf("hello");
}
'''
sim = MjSim(load_model_from_xml(XML.format(nuserdata=0)),
substep_callback=fn)
sim.step() # should print 'hello'
def test_sim_save():
model = load_model_from_xml(BASIC_MODEL_XML)
assert model.nkey == 0
sim = MjSim(model)
with StringIO() as f:
sim.save(f)
f.seek(0)
loaded_model = load_model_from_xml(f.read())
assert loaded_model.nkey == 1
with BytesIO() as f:
sim.save(f, format='mjb')
f.seek(0)
loaded_model = load_model_from_mjb(f.read())
assert loaded_model.nkey == 1
def _re_init(self, xml):
self.model = mujoco_py.load_model_from_xml(xml)
self.sim = mujoco_py.MjSim(self.model)
self.data = self.sim.data
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
if self.viewer:
self.viewer.update_sim(self.sim)
def main():
xml = None
with open(MODEL_XML_PATH, "r") as f:
xml = f.read()
model = load_model_from_xml(xml)
sim = MjSim(model)
nq = sim.data.qpos.shape[0]
nt = sim.data.ctrl.shape[0]
net, optimizer, loss_fn = gen_nn(nq, nt)
# load_nn_if_avaiable(net, optimizer)
tau_finals, qpos_finals = load_data()
iters = len(qpos_finals) // BATCH_SIZE
max_loss = []
max_var = []
avg_loss = []
for i in range(iters):
qpos_final_tensor = T.Tensor(
qpos_finals[:(i + 1) * BATCH_SIZE]).double().to(DEVICE)
tau_finals_tensor = T.Tensor(
tau_finals[:(i + 1) * BATCH_SIZE]).double().to(DEVICE)
tau_predicted = net(qpos_final_tensor).to(DEVICE)
optimizer.zero_grad()
loss = loss_fn(tau_predicted, tau_finals_tensor)
# Take loss statistics for mean and variance of current batch
max_loss.append(loss.max().item())
max_var.append(loss.var(dim=1).max().item())
# Calculate actual loss and backprop
loss = loss.mean()
avg_loss.append(loss.item())
print(f"Loss after {(i+1) * BATCH_SIZE} samples: {loss}")
loss.backward()
optimizer.step()
if i > 0 and i % BATCH_SIZE == 0:
print(f"Saving checkpoint and model info at iteration {i}")
np.savetxt(AVG_LOSS_PATH, np.array(avg_loss))
np.savetxt(MAX_LOSS_PATH, np.array(max_loss))
np.savetxt(MAX_VAR_PATH, np.array(max_var))
T.save(
{
'model_state_dict': net.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
}, NN_MODEL_CHECKPOINT_PATH)
T.save(net, NN_MODEL_PATH)
def mp_test_cameras():
model = load_model_from_xml(BASIC_MODEL_XML)
pool = MjRenderPool(model, n_workers=1)
image = pool.render(100, 100)
assert image.shape == (1, 100, 100, 3)
compare_imgs(image[0], 'test_render_pool.mp_test_cameras.0.png')
image = pool.render(100, 100, camera_name='camera1')
assert image.shape == (1, 100, 100, 3)
compare_imgs(image[0], 'test_render_pool.mp_test_cameras.1.png')
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_increment(self):
sim = MjSim(load_model_from_xml(XML.format(nuserdata=1)),
substep_callback=INCREMENT_FN)
self.assertEqual(sim.data.userdata[0], 0)
sim.step() # should increment userdata[0]
self.assertEqual(sim.data.userdata[0], 1)
# Test again with different nsubsteps, reuse model
sim = MjSim(sim.model, nsubsteps=7, substep_callback=INCREMENT_FN)
self.assertEqual(sim.data.userdata[0], 0)
sim.step() # should increment userdata[0] 7 times
self.assertEqual(sim.data.userdata[0], 7)
def get_sim(shape_infos, camera_info):
from mujoco_py import load_model_from_xml, MjSim
xml = """
<mujoco>
<!--
<default>
<geom material="DEFAULT_GEOM_MAT"/>
</default>
-->
<asset>
<texture name="SKYBOX" type="skybox" builtin="gradient" width="128" height="128" rgb1="0.4 0.6 0.8" rgb2="0 0 0"/>
<texture name="DEFAULT_GEOM_TEX" type="cube" builtin="flat" mark="cross" width="64" height="64" rgb1="0.8 0.6 0.4" rgb2="0.8 0.6 0.4" markrgb="1 1 1" random="0.01"/>
<material name="DEFAULT_GEOM_MAT" texture="DEFAULT_GEOM_TEX" texuniform="false"/>
</asset>
<worldbody>
<light diffuse=".5 .5 .5" pos="0 0 3" dir="0 0 -1"/>
<camera name="main" pos="{0} {1} {2}" fovy="45" mode="targetbody" target="plane"/>
<body name="plane">
<geom type="box" size="40 40 0.5" rgba=".9 0 0 1" pos="0 0 -0.5"/>
</body>
""".format(camera_info[0], camera_info[1], camera_info[2])
for i, shape_info in enumerate(shape_infos):
tx = shape_info['tx']
ty = shape_info['ty']
tz = shape_info['tz']
sx = shape_info['sx']
sy = shape_info['sy']
sz = shape_info["sz"]
r, g, b, a = shape_info['rgba']
type_cube = shape_info['type']
xml += """
<body name="cube_{6}">
<joint axis="1 0 0" type="hinge"/>
<joint axis="0 1 0" type="hinge"/>
<joint axis="0 0 1" type="hinge"/>
<geom pos="{0} {1} {2}" size="{3} {4} {5}" type="{7}" material="DEFAULT_GEOM_MAT" rgba="{8} {9} {10} {11}"/>
</body>
""".format(tx, ty, tz, sx, sy, sz, i, type_cube, r, g, b, a)
xml += """
</worldbody>
</mujoco>
"""
model = load_model_from_xml(xml)
sim = MjSim(model)
return sim
def test_warm_start_ema():
"""
Test that the smoothed commanded position is initialized to the commanded position at
start.
"""
# Load model with EMA smoothing (0.97)
model = load_model_from_xml(MODEL_XML.format(actuator=CASCADED_PDPI_ACTUATOR, nuser_actuator=1))
sim = MjSim(model)
cymj.set_pid_control(sim.model, sim.data)
# Load model without EMA smoothing (0.0)
model2 = load_model_from_xml(MODEL_XML.format(actuator=CASCADED_PDPI_ACTUATOR_NO_EMA, nuser_actuator=1))
sim2 = MjSim(model2)
cymj.set_pid_control(sim2.model, sim2.data)
init_pos = 0.1 * (random.random() - 0.5)
sim.data.ctrl[0] = sim2.data.ctrl[0] = .2
sim.step()
sim2.step()
EMA_SMOOTH_IDX = 4
# Compare that the saved smoothed position is the same for both EMA=.97 and EMA=0
assert abs(sim.data.userdata[EMA_SMOOTH_IDX] - sim2.data.userdata[EMA_SMOOTH_IDX]) < 1e-10
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_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 main():
xml = None
with open(MODEL_XML_PATH, "r") as xml_file:
xml = xml_file.read()
model = load_model_from_xml(xml)
sim = MjSim(model)
# viewer = MjViewer(sim)
# Basically numpy is dumb when loading 1D data
# To avoid enumerate(.) issues later on in code,
# make sure that we turn these into 2D arrays with
# a single column (or row, I don't know)
tau_finals, qpos_finals = load_data_if_available()
assert len(tau_finals) == len(qpos_finals)
tau_finals = np.array(tau_finals)
qpos_finals = np.array(qpos_finals)
if len(qpos_finals.shape) == 1:
qpos_finals.shape = (-1, 1)
if len(tau_finals.shape) == 1:
tau_finals.shape = (-1, 1)
num_samples = len(tau_finals)
good_q_samples = []
good_t_samples = []
for i in range(num_samples):
if i % 1000 == 0:
print(f"Currently running step {i}")
qpos = qpos_finals[i]
tau = tau_finals[i]
reset_sim(sim, qpos=qpos, taus=tau)
for _ in range(SIM_STEPS_TO_TEST):
sim.step()
if not check_end_of_sim(sim) or check_collision(sim):
break
if not check_end_of_sim(sim) or check_bad_gravity_point(sim):
continue # We have a bad sample
good_q_samples.append(qpos)
good_t_samples.append(tau)
print(f"Number of samples: {num_samples}")
print(f"Number of good samples: {len(good_q_samples)}")
assert len(good_q_samples) == len(good_t_samples)
np.savetxt(QPOS_DATA_PATH, np.array(good_q_samples))
np.savetxt(TAU_DATA_PATH, np.array(good_t_samples))
def test_mj_proportional_only():
"""
Check new PID control is backward compatible with position control
when it only has a Kp term.
"""
model = load_model_from_xml(MODEL_XML.format(actuator=P_ONLY_ACTUATOR, nuser_actuator=1))
sim = MjSim(model)
cymj.set_pid_control(sim.model, sim.data)
model2 = load_model_from_xml(MODEL_XML.format(actuator=POSITION_ACTUATOR, nuser_actuator=1))
sim2 = MjSim(model2)
init_pos = 0.1 * (random.random() - 0.5)
sim.data.qpos[0] = sim2.data.qpos[0] = init_pos
sim.data.ctrl[0] = sim2.data.ctrl[0] = 0
for i in range(2000):
print(i, sim.data.qpos[0], sim2.data.qpos[0])
sim.step()
sim2.step()
assert abs(sim.data.qpos[0] - sim2.data.qpos[0]) <= 1e-7, "%d step violates" % i
def initialize(self, use_cur_pos):
tmp = MODEL_XML_BASE.format(self.get_asset_mesh_str(),
self.get_asset_material_str(),
self.get_body_str())
model = load_model_from_xml(tmp)
self.sim = MjSim(model)
if self.view:
self.viewer = MjViewer(self.sim)
else:
self.viewer = mujoco_py.MjRenderContextOffscreen(self.sim, -1)
self._get_starting_step(use_cur_pos)
def test_mj_sim_buffers():
model = load_model_from_xml(BASIC_MODEL_XML)
# test no callback
sim = MjSim(model, nsubsteps=2)
assert(sim.udd_state == {})
sim.step()
assert(sim.udd_state == {})
# test with callback
foo = 10
d = {"foo": foo,
"foo_2": np.array([foo, foo])}
def udd_callback(sim):
return d
sim = MjSim(model, nsubsteps=2, udd_callback=udd_callback)
assert(sim.udd_state is not None)
assert(sim.udd_state["foo"] == foo)
assert(sim.udd_state["foo_2"].shape[0] == 2)
assert(sim.udd_state["foo_2"][0] == foo)
foo = 11
d = {"foo": foo,
"foo_2": np.array([foo, foo])}
sim.step()
assert(sim.udd_state is not None)
assert(sim.udd_state["foo"] == foo)
assert(sim.udd_state["foo_2"][0] == foo)
d = {}
with pytest.raises(AssertionError):
sim.step()
d = {"foo": foo,
"foo_2": np.array([foo, foo]),
"foo_3": foo}
with pytest.raises(AssertionError):
sim.step()
d = {"foo": foo,
"foo_2": np.array([foo, foo, foo])}
with pytest.raises(AssertionError):
sim.step()
d = {"foo": "haha",
"foo_2": np.array([foo, foo, foo])}
with pytest.raises(AssertionError):
sim.step()
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_mj_sim_buffers():
model = load_model_from_xml(BASIC_MODEL_XML)
# test no callback
sim = MjSim(model, nsubsteps=2)
assert(sim.udd_state == {})
sim.step()
assert(sim.udd_state == {})
# test with callback
foo = 10
d = {"foo": foo,
"foo_2": np.array([foo, foo])}
def udd_callback(sim):
return d
sim = MjSim(model, nsubsteps=2, udd_callback=udd_callback)
assert(sim.udd_state is not None)
assert(sim.udd_state["foo"] == foo)
assert(sim.udd_state["foo_2"].shape[0] == 2)
assert(sim.udd_state["foo_2"][0] == foo)
foo = 11
d = {"foo": foo,
"foo_2": np.array([foo, foo])}
sim.step()
assert(sim.udd_state is not None)
assert(sim.udd_state["foo"] == foo)
assert(sim.udd_state["foo_2"][0] == foo)
d = {}
with pytest.raises(AssertionError):
sim.step()
d = {"foo": foo,
"foo_2": np.array([foo, foo]),
"foo_3": foo}
with pytest.raises(AssertionError):
sim.step()
d = {"foo": foo,
"foo_2": np.array([foo, foo, foo])}
with pytest.raises(AssertionError):
sim.step()
d = {"foo": "haha",
"foo_2": np.array([foo, foo, foo])}
with pytest.raises(AssertionError):
sim.step()
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 _create_mujoco_model(self):
"""
Called to update the MuJoCo model by rewriting and reloading the XML file.
.. note::
This function is called from the constructor and from the domain parameter setter.
"""
xml_model = self.xml_model_template # don't change the template
xml_model = self._adapt_model_file(xml_model, self.domain_param)
# Create MuJoCo model from parsed XML file
self.model = mujoco_py.load_model_from_xml(xml_model)
self.sim = mujoco_py.MjSim(self.model, nsubsteps=self.frame_skip)
def mp_test_rendering():
model = load_model_from_xml(BASIC_MODEL_XML)
pool = MjRenderPool(model, n_workers=3)
images = pool.render(100, 100)
assert images.shape == (3, 100, 100, 3)
compare_imgs(images[0], 'test_render_pool.mp_test_rendering.0.png')
assert np.all(images[0] == images[1])
images, depth = pool.render(101, 103, depth=True)
assert images.shape == (3, 103, 101, 3)
assert depth.shape == (3, 103, 101)
assert np.all(images[0] == images[1])
assert np.all(images[1] == images[2])
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_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_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_mj_sim_pool_buffers():
model = load_model_from_xml(BASIC_MODEL_XML)
foo = 10
def udd_callback(sim):
return {"foo": foo}
sims = [MjSim(model, udd_callback=udd_callback) for _ in range(2)]
sim_pool = MjSimPool(sims, nsubsteps=2)
for i in range(len(sim_pool.sims)):
assert(sim_pool.sims[i].udd_state is not None)
assert(sim_pool.sims[i].udd_state["foo"] == 10)
foo = 11
sim_pool.step()
for i in range(len(sim_pool.sims)):
assert(sim_pool.sims[i].udd_state is not None)
assert(sim_pool.sims[i].udd_state["foo"] == 11)
def test_mj_warning_raises():
''' Test that MuJoCo warnings cause exceptions. '''
# Two boxes on a plane need more than 1 contact (nconmax)
xml = '''
<mujoco>
<size nconmax="1"/>
<worldbody>
<geom type="plane" size="1 1 0.1"/>
<body pos="1 0 1"> <joint type="free"/> <geom size="1"/> </body>
<body pos="0 1 1"> <joint type="free"/> <geom size="1"/> </body>
</worldbody>
</mujoco>
'''
model = load_model_from_xml(xml)
sim = MjSim(model)
sim.reset()
with pytest.raises(Exception):
# This should raise an exception due to the mujoco warning callback
sim.step()
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_sensors():
model = load_model_from_xml(BASIC_MODEL_XML)
sim = MjSim(model)
sim.model.sensor_names
sim.data.get_sensor("touchsensor")
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_load_particle():
load_model_from_xml(PARTICLE_XML)