From cef5b31a3fb83490ab70d1da87880633f1507a87 Mon Sep 17 00:00:00 2001 From: Alessio Quaglino Date: Tue, 11 Feb 2025 05:49:48 -0800 Subject: [PATCH] Drop `find_` prefix in Python MjSpec methods accessing elements by name. PiperOrigin-RevId: 725594833 Change-Id: Ibf8a9918e98b298751d0ef2e657714e28e5fa460 --- doc/python.rst | 5 +-- python/mjspec.ipynb | 38 ++++++++++--------- .../mujoco/codegen/generate_spec_bindings.py | 2 +- python/mujoco/specs_test.py | 24 ++++++------ 4 files changed, 35 insertions(+), 34 deletions(-) diff --git a/doc/python.rst b/doc/python.rst index a77306b807..8f8f520e13 100644 --- a/doc/python.rst +++ b/doc/python.rst @@ -589,9 +589,8 @@ to make model editing easier: Finding elements ^^^^^^^^^^^^^^^^ -The ``MjSpec`` object has methods like ``find_body(), find_joint(), find_site(), ...`` for every type of name-able -:ref:`mjtObj` type. These methods take a name string and return the element with that name or ``None`` if it does not -exist. +The ``MjSpec`` object has methods like ``.body(), .joint(), .site(), ...`` for named access of elements. +``spec.geom('my_geom')`` will return the :ref:`mjsGeom` called "my_geom", or ``None`` if it does not exist. Element lists ^^^^^^^^^^^^^ diff --git a/python/mjspec.ipynb b/python/mjspec.ipynb index d813ea9537..e5392ed2f9 100644 --- a/python/mjspec.ipynb +++ b/python/mjspec.ipynb @@ -573,11 +573,11 @@ " self.spec = mj.MjSpec.from_string(leg_model)\n", "\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", "\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.pos[0] = length" ] @@ -620,6 +620,7 @@ " \"\"\"Constructs a creature with `num_legs` legs.\"\"\"\n", " rgba = random_state.uniform([0, 0, 0, 1], [1, 1, 1, 1])\n", " spec = mj.MjSpec.from_string(creature_model)\n", + " spec.copy_during_attach = True\n", "\n", " # Attach legs to equidistant sites on the circumference.\n", " spec.worldbody.first_geom().rgba = rgba\n", @@ -628,7 +629,7 @@ " theta = 2 * i * np.pi / num_legs\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.attach_body(leg.spec.find_body('thigh'), '', '-' + str(i))\n", + " hip_site.attach_body(leg.spec.body('thigh'), '', '-' + str(i))\n", "\n", " return spec" ] @@ -870,7 +871,7 @@ "# # Add another humanoid\n", "spec_humanoid = mj.MjSpec.from_file(humanoid_file)\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", "# Recompile preserving the state\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", "\n", "spec = mj.MjSpec.from_file(humanoid_file)\n", + "spec.copy_during_attach = True\n", "\n", "# Get the torso, arm, and leg bodies\n", - "arm_left = spec.find_body('upper_arm_left')\n", - "arm_right = spec.find_body('upper_arm_right')\n", - "leg_left = spec.find_body('thigh_left')\n", - "leg_right = spec.find_body('thigh_right')\n", - "torso = spec.find_body('torso')\n", + "arm_left = spec.body('upper_arm_left')\n", + "arm_right = spec.body('upper_arm_right')\n", + "leg_left = spec.body('thigh_left')\n", + "leg_right = spec.body('thigh_right')\n", + "torso = spec.body('torso')\n", "\n", "# Attach frames at the arm positions\n", "shoulder_left = torso.add_frame(pos=arm_left.pos)\n", @@ -976,13 +978,13 @@ " spec.degree = False # MuJoCo release\n", "\n", "# Replace right arm with frame\n", - "arm_right = spec.find_body('upper_arm_right')\n", - "torso = spec.find_body('torso')\n", + "arm_right = spec.body('upper_arm_right')\n", + "torso = spec.body('torso')\n", "shoulder_right = torso.add_frame(pos=arm_right.pos, quat=[0, 0.8509035, 0, 0.525322])\n", "spec.detach_body(arm_right)\n", "\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", "\n", "model = spec.compile()\n", @@ -1033,6 +1035,7 @@ "\n", "humanoid = mj.MjSpec.from_file(humanoid_file)\n", "spec = mj.MjSpec()\n", + "spec.copy_during_attach = True\n", "\n", "# Delete all key frames to avoid name conflicts\n", "while humanoid.keys:\n", @@ -1044,18 +1047,18 @@ " humanoid.materials[0].rgba = [\n", " np.random.uniform(), np.random.uniform(),\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", - " 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()] # 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()] # Randomize right arm orientation\n", "\n", " # attach randomized humanoid to parent spec\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", "spec.worldbody.add_light(mode=mj.mjtCamLight.mjCAMLIGHT_TARGETBODYCOM,\n", " targetbody='3torso3', diffuse=[.8, .8, .8],\n", @@ -1072,8 +1075,7 @@ "sJFuNetilv4m" ], "gpuClass": "premium", - "private_outputs": true, - "toc_visible": true + "private_outputs": true }, "gpuClass": "premium", "kernelspec": { diff --git a/python/mujoco/codegen/generate_spec_bindings.py b/python/mujoco/codegen/generate_spec_bindings.py index 705cd8169a..1ee86ce94a 100644 --- a/python/mujoco/codegen/generate_spec_bindings.py +++ b/python/mujoco/codegen/generate_spec_bindings.py @@ -583,7 +583,7 @@ def generate_find() -> None: elemlower = elem.lower() titlecase = 'Mjs' + elem code = f"""\n - mjSpec.def("find_{elemlower}", + mjSpec.def("{elemlower}", [](MjSpec& self, std::string& name) -> raw::{titlecase}* {{ return mjs_as{elem}( mjs_findElement(self.ptr, {objtype}, name.c_str())); diff --git a/python/mujoco/specs_test.py b/python/mujoco/specs_test.py index d7eb9b0c9f..470943b01d 100644 --- a/python/mujoco/specs_test.py +++ b/python/mujoco/specs_test.py @@ -617,12 +617,12 @@ def test_element_list(self): self.assertEqual(spec.actuators[0].name, 'actuator1') self.assertEqual(spec.actuators[1].name, 'actuator2') self.assertEqual(spec.actuators[2].name, 'actuator3') - self.assertEqual(spec.find_sensor('sensor1'), sensor1) - self.assertEqual(spec.find_sensor('sensor2'), sensor2) - self.assertEqual(spec.find_sensor('sensor3'), sensor3) - self.assertEqual(spec.find_actuator('actuator1'), actuator1) - self.assertEqual(spec.find_actuator('actuator2'), actuator2) - self.assertEqual(spec.find_actuator('actuator3'), actuator3) + self.assertEqual(spec.sensor('sensor1'), sensor1) + self.assertEqual(spec.sensor('sensor2'), sensor2) + self.assertEqual(spec.sensor('sensor3'), sensor3) + self.assertEqual(spec.actuator('actuator1'), actuator1) + self.assertEqual(spec.actuator('actuator2'), actuator2) + self.assertEqual(spec.actuator('actuator3'), actuator3) def test_body_list(self): main_xml = """ @@ -709,12 +709,12 @@ def test_geom_list(self): spec = mujoco.MjSpec.from_string(main_xml) geom1 = spec.worldbody.add_geom(name='geom1') 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.find_geom('geom1'), geom1) - self.assertEqual(spec.find_geom('geom2'), geom2) - self.assertEqual(spec.find_geom('geom3'), geom3) + self.assertEqual(spec.geom('geom1'), geom1) + self.assertEqual(spec.geom('geom2'), geom2) + self.assertEqual(spec.geom('geom3'), geom3) def test_iterators(self): spec = mujoco.MjSpec() @@ -801,11 +801,11 @@ def test_delete(self): self.assertEqual(model.nsite, 11) self.assertEqual(model.nsensor, 11) - head = spec.find_body('head') + head = spec.body('head') self.assertIsNotNone(head) site = head.first_site() self.assertIsNotNone(site) - self.assertEqual(site, spec.find_site('head')) + self.assertEqual(site, spec.site('head')) site.delete() spec.sensors[-1].delete()