"""Tests for sim.py.""" import mujoco import mujoco_warp as mjwarp import numpy as np import pytest import torch from conftest import get_test_device from mjlab.sim import MujocoCfg, Simulation, SimulationCfg @pytest.fixture def device(): """Test device fixture.""" return get_test_device() @pytest.fixture def robot_xml(): """Simple robot with geoms and joints.""" return """ """ def test_simulation_config_is_piped(robot_xml, device): """Test that SimulationCfg values are applied to both mj_model and wp_model.""" model = mujoco.MjModel.from_xml_string(robot_xml) cfg = SimulationCfg( contact_sensor_maxmatch=128, broadphase="sap_tile", broadphase_filter=("plane", "aabb"), mujoco=MujocoCfg( timestep=0.02, integrator="euler", solver="cg", iterations=7, ls_iterations=14, ccd_iterations=20, gravity=(0, 0, 7.5), enableflags=("energy",), ), ) sim = Simulation(num_envs=1, cfg=cfg, model=model, device=device) # MujocoCfg should be applied to mj_model. assert sim.mj_model.opt.timestep == cfg.mujoco.timestep assert sim.mj_model.opt.integrator == mujoco.mjtIntegrator.mjINT_EULER assert sim.mj_model.opt.solver == mujoco.mjtSolver.mjSOL_CG assert sim.mj_model.opt.iterations == cfg.mujoco.iterations assert sim.mj_model.opt.ls_iterations == cfg.mujoco.ls_iterations assert sim.mj_model.opt.ccd_iterations == cfg.mujoco.ccd_iterations assert tuple(sim.mj_model.opt.gravity) == cfg.mujoco.gravity assert sim.mj_model.opt.enableflags & mujoco.mjtEnableBit.mjENBL_ENERGY # MujocoCfg should be inherited by wp_model via put_model. np.testing.assert_almost_equal( sim.model.opt.timestep[0].cpu().numpy(), cfg.mujoco.timestep ) np.testing.assert_almost_equal( sim.model.opt.gravity[0].cpu().numpy(), cfg.mujoco.gravity ) assert sim.model.opt.integrator == mujoco.mjtIntegrator.mjINT_EULER assert sim.model.opt.solver == mujoco.mjtSolver.mjSOL_CG assert sim.model.opt.iterations == cfg.mujoco.iterations assert sim.model.opt.enableflags & mujoco.mjtEnableBit.mjENBL_ENERGY # SimulationCfg's warp-only settings should be applied to wp_model.opt. assert sim.wp_model.opt.contact_sensor_maxmatch == cfg.contact_sensor_maxmatch assert sim.wp_model.opt.broadphase == mjwarp.BroadphaseType.SAP_TILE assert sim.wp_model.opt.broadphase_filter == ( mjwarp.BroadphaseFilter.PLANE | mjwarp.BroadphaseFilter.AABB ) def test_default_broadphase_keeps_put_model_heuristic(robot_xml, device): """Unset broadphase settings should not override put_model's own heuristic.""" model = mujoco.MjModel.from_xml_string(robot_xml) heuristic_opt = mjwarp.put_model(model).opt sim = Simulation(num_envs=1, cfg=SimulationCfg(), model=model, device=device) assert sim.wp_model.opt.broadphase == heuristic_opt.broadphase assert sim.wp_model.opt.broadphase_filter == heuristic_opt.broadphase_filter def test_ls_parallel_is_deprecated(): """Setting the removed ls_parallel option warns instead of erroring.""" with pytest.warns(DeprecationWarning, match="ls_parallel"): SimulationCfg(ls_parallel=True) def test_sim_reset_restores_initial_state(robot_xml, device): """Test that sim.reset() restores qpos/qvel to initial values.""" model = mujoco.MjModel.from_xml_string(robot_xml) sim = Simulation(num_envs=2, cfg=SimulationCfg(), model=model, device=device) qpos0 = sim.data.qpos.clone() qvel0 = sim.data.qvel.clone() # Run simulation to modify state. for _ in range(10): sim.step() assert not torch.allclose(sim.data.qpos, qpos0) assert not torch.allclose(sim.data.qvel, qvel0) # Reset should restore initial state. sim.reset() torch.testing.assert_close(sim.data.qpos[:], qpos0) torch.testing.assert_close(sim.data.qvel[:], qvel0) # qacc_warmstart should be zeroed. assert (sim.data.qacc_warmstart == 0).all() @pytest.mark.skipif(not torch.cuda.is_available(), reason="Likely bug on CPU MjWarp") def test_sim_reset_selective(robot_xml, device): """Test that sim.reset() only affects specified environments.""" model = mujoco.MjModel.from_xml_string(robot_xml) sim = Simulation(num_envs=4, cfg=SimulationCfg(), model=model, device=device) qpos0 = sim.data.qpos.clone() # Run simulation to modify state. for _ in range(10): sim.step() qpos_after_sim = sim.data.qpos.clone() # Reset only env 1 and 3. sim.reset(torch.tensor([1, 3], device=device)) # Envs 1 and 3 should be reset. torch.testing.assert_close(sim.data.qpos[1], qpos0[1]) torch.testing.assert_close(sim.data.qpos[3], qpos0[3]) # Envs 0 and 2 should be unchanged. torch.testing.assert_close(sim.data.qpos[0], qpos_after_sim[0]) torch.testing.assert_close(sim.data.qpos[2], qpos_after_sim[2]) def test_xpos_matches_qpos_after_forward(robot_xml, device): """sim.step() leaves xpos stale; sim.forward() makes it match qpos. In MuJoCo, mj_step = mj_step1 (forward kinematics + forces) + mj_step2 (integration). After mj_step, qpos/qvel are post-integration but xpos is from the pre-integration forward pass. sim.forward() recomputes xpos from the current qpos. """ model = mujoco.MjModel.from_xml_string(robot_xml) cfg = SimulationCfg(mujoco=MujocoCfg(timestep=0.01)) # Large dt for clear signal sim = Simulation(num_envs=2, cfg=cfg, model=model, device=device) # Step enough for significant velocity -> large staleness gap. for _ in range(50): sim.step() # xpos is stale: reflects pre-integration state of last step. # For the freejoint body (body 1), qpos[:3] is the true position. xpos_stale = sim.data.xpos[:, 1].clone() qpos_pos = sim.data.qpos[:, :3].clone() assert not torch.allclose(xpos_stale, qpos_pos, atol=1e-4) # forward() refreshes derived quantities from current qpos. sim.forward() xpos_fresh = sim.data.xpos[:, 1].clone() torch.testing.assert_close(xpos_fresh, qpos_pos, atol=1e-5, rtol=0)