Spaces:
Paused
Paused
File size: 3,761 Bytes
225dc91 7f9f21f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | 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
|