chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Drop find_ prefix in Python MjSpec methods accessing elements by name.

Browse Source 
PiperOrigin-RevId: 725594833
Change-Id: Ibf8a9918e98b298751d0ef2e657714e28e5fa460
This commit is contained in:
Alessio Quaglino
2025-02-11 05:49:48 -08:00
committed by Copybara-Service
parent a6dd5f086f
commit cef5b31a3f
4 changed files with 35 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/python.rst
+2 -3
View File
@@ -589,9 +589,8 @@ to make model editing easier:
Finding elements Finding elements
^^^^^^^^^^^^^^^^ ^^^^^^^^^^^^^^^^
The ``MjSpec`` object has methods like ``find_body(), find_joint(), find_site(), ...`` for every type of name-able The ``MjSpec`` object has methods like ``.body(), .joint(), .site(), ...`` for named access of elements.
:ref:`mjtObj` type. These methods take a name string and return the element with that name or ``None`` if it does not ``spec.geom('my_geom')`` will return the :ref:`mjsGeom` called "my_geom", or ``None`` if it does not exist.
exist.
Element lists Element lists
^^^^^^^^^^^^^ ^^^^^^^^^^^^^
python/mjspec.ipynb
+20 -18
View File
@@ -573,11 +573,11 @@
" self.spec = mj.MjSpec.from_string(leg_model)\n", " self.spec = mj.MjSpec.from_string(leg_model)\n",
"\n", "\n",
" # Thigh:\n", " # Thigh:\n",
" thigh = self.spec.find_body('thigh')\n", " thigh = self.spec.body('thigh')\n",
" thigh.add_geom(fromto=[0, 0, 0, length, 0, 0], size=[length/4, 0, 0], rgba=rgba)\n", " thigh.add_geom(fromto=[0, 0, 0, length, 0, 0], size=[length/4, 0, 0], rgba=rgba)\n",
"\n", "\n",
" # Hip:\n", " # Hip:\n",
" shin = self.spec.find_body('shin')\n", " shin = self.spec.body('shin')\n",
" shin.add_geom(fromto=[0, 0, 0, 0, 0, -length], size=[length/5, 0, 0], rgba=rgba)\n", " shin.add_geom(fromto=[0, 0, 0, 0, 0, -length], size=[length/5, 0, 0], rgba=rgba)\n",
" shin.pos[0] = length" " shin.pos[0] = length"
] ]
@@ -620,6 +620,7 @@
" \"\"\"Constructs a creature with `num_legs` legs.\"\"\"\n", " \"\"\"Constructs a creature with `num_legs` legs.\"\"\"\n",
" rgba = random_state.uniform([0, 0, 0, 1], [1, 1, 1, 1])\n", " rgba = random_state.uniform([0, 0, 0, 1], [1, 1, 1, 1])\n",
" spec = mj.MjSpec.from_string(creature_model)\n", " spec = mj.MjSpec.from_string(creature_model)\n",
" spec.copy_during_attach = True\n",
"\n", "\n",
" # Attach legs to equidistant sites on the circumference.\n", " # Attach legs to equidistant sites on the circumference.\n",
" spec.worldbody.first_geom().rgba = rgba\n", " spec.worldbody.first_geom().rgba = rgba\n",
@@ -628,7 +629,7 @@
" theta = 2 * i * np.pi / num_legs\n", " theta = 2 * i * np.pi / num_legs\n",
" hip_pos = BODY_RADIUS * np.array([np.cos(theta), np.sin(theta), 0])\n", " hip_pos = BODY_RADIUS * np.array([np.cos(theta), np.sin(theta), 0])\n",
" hip_site = spec.worldbody.add_site(pos=hip_pos, euler=[0, 0, theta])\n", " hip_site = spec.worldbody.add_site(pos=hip_pos, euler=[0, 0, theta])\n",
" hip_site.attach_body(leg.spec.find_body('thigh'), '', '-' + str(i))\n", " hip_site.attach_body(leg.spec.body('thigh'), '', '-' + str(i))\n",
"\n", "\n",
" return spec" " return spec"
] ]
@@ -870,7 +871,7 @@
"# # Add another humanoid\n", "# # Add another humanoid\n",
"spec_humanoid = mj.MjSpec.from_file(humanoid_file)\n", "spec_humanoid = mj.MjSpec.from_file(humanoid_file)\n",
"attachment_frame = spec.worldbody.add_frame(pos=[0, -1, 2])\n", "attachment_frame = spec.worldbody.add_frame(pos=[0, -1, 2])\n",
"attachment_frame.attach_body(spec_humanoid.find_body('torso'), 'a', 'b')\n", "attachment_frame.attach_body(spec_humanoid.body('torso'), 'a', 'b')\n",
"\n", "\n",
"# Recompile preserving the state\n", "# Recompile preserving the state\n",
"new_model, new_data = spec.recompile(model, data)\n", "new_model, new_data = spec.recompile(model, data)\n",
@@ -924,13 +925,14 @@
"#@title Humanoid with arms replaced by legs.{vertical-output: true}\n", "#@title Humanoid with arms replaced by legs.{vertical-output: true}\n",
"\n", "\n",
"spec = mj.MjSpec.from_file(humanoid_file)\n", "spec = mj.MjSpec.from_file(humanoid_file)\n",
"spec.copy_during_attach = True\n",
"\n", "\n",
"# Get the torso, arm, and leg bodies\n", "# Get the torso, arm, and leg bodies\n",
"arm_left = spec.find_body('upper_arm_left')\n", "arm_left = spec.body('upper_arm_left')\n",
"arm_right = spec.find_body('upper_arm_right')\n", "arm_right = spec.body('upper_arm_right')\n",
"leg_left = spec.find_body('thigh_left')\n", "leg_left = spec.body('thigh_left')\n",
"leg_right = spec.find_body('thigh_right')\n", "leg_right = spec.body('thigh_right')\n",
"torso = spec.find_body('torso')\n", "torso = spec.body('torso')\n",
"\n", "\n",
"# Attach frames at the arm positions\n", "# Attach frames at the arm positions\n",
"shoulder_left = torso.add_frame(pos=arm_left.pos)\n", "shoulder_left = torso.add_frame(pos=arm_left.pos)\n",
@@ -976,13 +978,13 @@
" spec.degree = False # MuJoCo release\n", " spec.degree = False # MuJoCo release\n",
"\n", "\n",
"# Replace right arm with frame\n", "# Replace right arm with frame\n",
"arm_right = spec.find_body('upper_arm_right')\n", "arm_right = spec.body('upper_arm_right')\n",
"torso = spec.find_body('torso')\n", "torso = spec.body('torso')\n",
"shoulder_right = torso.add_frame(pos=arm_right.pos, quat=[0, 0.8509035, 0, 0.525322])\n", "shoulder_right = torso.add_frame(pos=arm_right.pos, quat=[0, 0.8509035, 0, 0.525322])\n",
"spec.detach_body(arm_right)\n", "spec.detach_body(arm_right)\n",
"\n", "\n",
"# Attach Franka arm to humanoid\n", "# Attach Franka arm to humanoid\n",
"franka_arm = franka.find_body('fr3_link2')\n", "franka_arm = franka.body('fr3_link2')\n",
"shoulder_right.attach_body(franka_arm, 'franka', '')\n", "shoulder_right.attach_body(franka_arm, 'franka', '')\n",
"\n", "\n",
"model = spec.compile()\n", "model = spec.compile()\n",
@@ -1033,6 +1035,7 @@
"\n", "\n",
"humanoid = mj.MjSpec.from_file(humanoid_file)\n", "humanoid = mj.MjSpec.from_file(humanoid_file)\n",
"spec = mj.MjSpec()\n", "spec = mj.MjSpec()\n",
"spec.copy_during_attach = True\n",
"\n", "\n",
"# Delete all key frames to avoid name conflicts\n", "# Delete all key frames to avoid name conflicts\n",
"while humanoid.keys:\n", "while humanoid.keys:\n",
@@ -1044,18 +1047,18 @@
" humanoid.materials[0].rgba = [\n", " humanoid.materials[0].rgba = [\n",
" np.random.uniform(), np.random.uniform(),\n", " np.random.uniform(), np.random.uniform(),\n",
" np.random.uniform(), 1] # Randomize color\n", " np.random.uniform(), 1] # Randomize color\n",
" humanoid.find_body('head').first_geom().size = [\n", " humanoid.body('head').first_geom().size = [\n",
" .18*np.random.uniform(), 0, 0] # Randomize head size\n", " .18*np.random.uniform(), 0, 0] # Randomize head size\n",
" humanoid.find_body('upper_arm_left').quat = [\n", " humanoid.body('upper_arm_left').quat = [\n",
" np.random.uniform(), np.random.uniform(),\n", " np.random.uniform(), np.random.uniform(),\n",
" np.random.uniform(), np.random.uniform()] # Randomize left arm orientation\n", " np.random.uniform(), np.random.uniform()] # Randomize left arm orientation\n",
" humanoid.find_body('upper_arm_right').quat = [\n", " humanoid.body('upper_arm_right').quat = [\n",
" np.random.uniform(), np.random.uniform(),\n", " np.random.uniform(), np.random.uniform(),\n",
" np.random.uniform(), np.random.uniform()] # Randomize right arm orientation\n", " np.random.uniform(), np.random.uniform()] # Randomize right arm orientation\n",
"\n", "\n",
" # attach randomized humanoid to parent spec\n", " # attach randomized humanoid to parent spec\n",
" frame = spec.worldbody.add_frame(pos=[i, j, 0])\n", " frame = spec.worldbody.add_frame(pos=[i, j, 0])\n",
" frame.attach_body(humanoid.find_body('torso'), str(i), str(j))\n", " frame.attach_body(humanoid.body('torso'), str(i), str(j))\n",
"\n", "\n",
"spec.worldbody.add_light(mode=mj.mjtCamLight.mjCAMLIGHT_TARGETBODYCOM,\n", "spec.worldbody.add_light(mode=mj.mjtCamLight.mjCAMLIGHT_TARGETBODYCOM,\n",
" targetbody='3torso3', diffuse=[.8, .8, .8],\n", " targetbody='3torso3', diffuse=[.8, .8, .8],\n",
@@ -1072,8 +1075,7 @@
"sJFuNetilv4m" "sJFuNetilv4m"
], ],
"gpuClass": "premium", "gpuClass": "premium",
"private_outputs": true, "private_outputs": true
"toc_visible": true
}, },
"gpuClass": "premium", "gpuClass": "premium",
"kernelspec": { "kernelspec": {
python/mujoco/codegen/generate_spec_bindings.py
+1 -1
View File
@@ -583,7 +583,7 @@ def generate_find() -> None:
elemlower = elem.lower() elemlower = elem.lower()
titlecase = 'Mjs' + elem titlecase = 'Mjs' + elem
code = f"""\n code = f"""\n
mjSpec.def("find_{elemlower}", mjSpec.def("{elemlower}",
[](MjSpec& self, std::string& name) -> raw::{titlecase}* {{ [](MjSpec& self, std::string& name) -> raw::{titlecase}* {{
return mjs_as{elem}( return mjs_as{elem}(
mjs_findElement(self.ptr, {objtype}, name.c_str())); mjs_findElement(self.ptr, {objtype}, name.c_str()));
python/mujoco/specs_test.py
+12 -12
View File
@@ -617,12 +617,12 @@ class SpecsTest(absltest.TestCase):
self.assertEqual(spec.actuators[0].name, 'actuator1') self.assertEqual(spec.actuators[0].name, 'actuator1')
self.assertEqual(spec.actuators[1].name, 'actuator2') self.assertEqual(spec.actuators[1].name, 'actuator2')
self.assertEqual(spec.actuators[2].name, 'actuator3') self.assertEqual(spec.actuators[2].name, 'actuator3')
self.assertEqual(spec.find_sensor('sensor1'), sensor1) self.assertEqual(spec.sensor('sensor1'), sensor1)
self.assertEqual(spec.find_sensor('sensor2'), sensor2) self.assertEqual(spec.sensor('sensor2'), sensor2)
self.assertEqual(spec.find_sensor('sensor3'), sensor3) self.assertEqual(spec.sensor('sensor3'), sensor3)
self.assertEqual(spec.find_actuator('actuator1'), actuator1) self.assertEqual(spec.actuator('actuator1'), actuator1)
self.assertEqual(spec.find_actuator('actuator2'), actuator2) self.assertEqual(spec.actuator('actuator2'), actuator2)
self.assertEqual(spec.find_actuator('actuator3'), actuator3) self.assertEqual(spec.actuator('actuator3'), actuator3)
def test_body_list(self): def test_body_list(self):
main_xml = """ main_xml = """
@@ -709,12 +709,12 @@ class SpecsTest(absltest.TestCase):
spec = mujoco.MjSpec.from_string(main_xml) spec = mujoco.MjSpec.from_string(main_xml)
geom1 = spec.worldbody.add_geom(name='geom1') geom1 = spec.worldbody.add_geom(name='geom1')
geom2 = spec.worldbody.add_geom(name='geom2') geom2 = spec.worldbody.add_geom(name='geom2')
geom3 = spec.find_body('body1').add_geom(name='geom3') geom3 = spec.body('body1').add_geom(name='geom3')
self.assertEqual(spec.geoms, [geom1, geom2, geom3]) self.assertEqual(spec.geoms, [geom1, geom2, geom3])
self.assertEqual(spec.find_geom('geom1'), geom1) self.assertEqual(spec.geom('geom1'), geom1)
self.assertEqual(spec.find_geom('geom2'), geom2) self.assertEqual(spec.geom('geom2'), geom2)
self.assertEqual(spec.find_geom('geom3'), geom3) self.assertEqual(spec.geom('geom3'), geom3)
def test_iterators(self): def test_iterators(self):
spec = mujoco.MjSpec() spec = mujoco.MjSpec()
@@ -801,11 +801,11 @@ class SpecsTest(absltest.TestCase):
self.assertEqual(model.nsite, 11) self.assertEqual(model.nsite, 11)
self.assertEqual(model.nsensor, 11) self.assertEqual(model.nsensor, 11)
head = spec.find_body('head') head = spec.body('head')
self.assertIsNotNone(head) self.assertIsNotNone(head)
site = head.first_site() site = head.first_site()
self.assertIsNotNone(site) self.assertIsNotNone(site)
self.assertEqual(site, spec.find_site('head')) self.assertEqual(site, spec.site('head'))
site.delete() site.delete()
spec.sensors[-1].delete() spec.sensors[-1].delete()