Skip to content

Commit

Permalink
Drop find_ prefix in Python MjSpec methods accessing elements by name.
Browse files Browse the repository at this point in the history
PiperOrigin-RevId: 725594833
Change-Id: Ibf8a9918e98b298751d0ef2e657714e28e5fa460
  • Loading branch information
quagla authored and copybara-github committed Feb 11, 2025
1 parent a6dd5f0 commit cef5b31
Show file tree
Hide file tree
Showing 4 changed files with 35 additions and 34 deletions.
5 changes: 2 additions & 3 deletions doc/python.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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
^^^^^^^^^^^^^
Expand Down
38 changes: 20 additions & 18 deletions python/mjspec.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -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"
]
Expand Down Expand Up @@ -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",
Expand All @@ -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"
]
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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",
Expand All @@ -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",
Expand All @@ -1072,8 +1075,7 @@
"sJFuNetilv4m"
],
"gpuClass": "premium",
"private_outputs": true,
"toc_visible": true
"private_outputs": true
},
"gpuClass": "premium",
"kernelspec": {
Expand Down
2 changes: 1 addition & 1 deletion python/mujoco/codegen/generate_spec_bindings.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()));
Expand Down
24 changes: 12 additions & 12 deletions python/mujoco/specs_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = """
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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()
Expand Down

0 comments on commit cef5b31

Please sign in to comment.