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)