Spaces:
Paused
Paused
| import mujoco | |
| import pytest | |
| from api.robots import ROBOT_NAMES, ROBOTS, reset_to_robot_keyframe | |
| from worker.sim_runner import RunConfig, _hold_controller, config_for, simulate | |
| # A hinged pole driven by a torque motor, starting 0.3 rad from upright: gravity pulls it over unless held. | |
| POLE = """<mujoco> | |
| <compiler angle="radian"/> | |
| <option timestep="0.002"/> | |
| <worldbody> | |
| <geom type="plane" size="1 1 0.1"/> | |
| <body name="base" pos="0 0 0.1"> | |
| <body name="pole"> | |
| <joint name="hinge" type="hinge" axis="0 1 0" range="-3 3"/> | |
| <geom type="capsule" fromto="0 0 0 0 0 1" size="0.03" mass="1"/> | |
| </body> | |
| </body> | |
| </worldbody> | |
| <actuator><motor joint="hinge" ctrlrange="-30 30"/></actuator> | |
| </mujoco>""" | |
| def pole_angle_after(tmp_path, controller, home=(0.3,)): | |
| (tmp_path / "scene.xml").write_text(POLE) | |
| config = RunConfig(2.0, controller, render=False, root_body="base", fall_height=None, home_qpos=home) | |
| simulate(tmp_path / "scene.xml", tmp_path / controller, config) | |
| model = mujoco.MjModel.from_xml_path(str(tmp_path / "scene.xml")) | |
| data = mujoco.MjData(model) | |
| reset_to_robot_keyframe(model, data, "base", home) | |
| hold = _hold_controller(model, data, config) | |
| if controller == "passive": | |
| model.opt.disableflags |= mujoco.mjtDisableBit.mjDSBL_ACTUATION | |
| for _ in range(1000): | |
| if controller == "hold": | |
| hold.apply(data) | |
| mujoco.mj_step(model, data) | |
| return float(data.qpos[0]) | |
| def test_hold_uses_pd_for_torque_motors(tmp_path): | |
| assert pole_angle_after(tmp_path, "hold") == pytest.approx(0.3, abs=0.05) | |
| assert abs(pole_angle_after(tmp_path, "passive")) > 1.0 # falls over without the controller | |
| def test_home_pose_is_used_only_without_a_keyframe(tmp_path): | |
| model = mujoco.MjModel.from_xml_string(POLE) | |
| data = mujoco.MjData(model) | |
| reset_to_robot_keyframe(model, data, "base", (0.3,)) | |
| assert data.qpos[0] == pytest.approx(0.3) | |
| keyed = mujoco.MjModel.from_xml_string( | |
| POLE.replace("</mujoco>", '<keyframe><key qpos="-0.2"/></keyframe></mujoco>') | |
| ) | |
| data = mujoco.MjData(keyed) | |
| reset_to_robot_keyframe(keyed, data, "base", (0.3,)) | |
| assert data.qpos[0] == pytest.approx(-0.2) | |
| def test_robot_configs(): | |
| assert set(ROBOT_NAMES) == {"unitree_g1", "unitree_h1", "unitree_go2", "franka_emika_panda"} | |
| arm = config_for("franka_emika_panda", 5, "hold") | |
| assert arm.root_body == "link0" and arm.fall_height is None and arm.camera_distance < 3 | |
| h1 = config_for("unitree_h1", 5, "hold") | |
| assert h1.home_qpos == ROBOTS["unitree_h1"].home_qpos and h1.kp == 1000 | |
| assert config_for(None, 5, "hold").root_body is None | |
| def test_api_accepts_every_robot(run_env): | |
| client, keys_, *_ = run_env | |
| key = keys_["alice"] | |
| r = client.post( | |
| "/simulate/generate", | |
| json={"description": "a quadruped near a box", "robot": "atlas"}, | |
| headers={"X-API-Key": key}, | |
| ) | |
| assert r.status_code == 422 and "unitree_go2" in r.text | |
| def test_robot_facts_with_unnamed_free_joint(tmp_path): | |
| """Go2 and H1 have an unnamed free joint; the facts given to the model must still build (was a TypeError).""" | |
| from api.robots import robot_facts | |
| spec = ROBOTS["unitree_go2"] | |
| (tmp_path / spec.directory).mkdir() | |
| (tmp_path / spec.directory / spec.model_file).write_text( | |
| '<mujoco><compiler angle="radian"/><worldbody><body name="base"><freejoint/><geom size=".1"/>' | |
| '<body name="leg"><joint name="knee" range="-1 1"/><geom size=".05"/></body></body></worldbody>' | |
| '<actuator><motor joint="knee"/></actuator></mujoco>' | |
| ) | |
| facts = robot_facts(str(tmp_path), "unitree_go2") | |
| assert "(unnamed joint 0)" in facts and "knee [-1.00, 1.00]" in facts | |