Skip to content

Add named element access methods to mjx.Model (body, joint, geom, actuator, etc.) #3002

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Open
Yuvraj-cyborg wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Add named element access methods to mjx.Model (body, joint, geom, actuator, etc.) #3002

Yuvraj-cyborg wants to merge 2 commits into from
+532 −1

Conversation

@Yuvraj-cyborg
Copy link

@Yuvraj-cyborg Yuvraj-cyborg commented Dec 29, 2025

Fixes: #2806

Problem :

Users of mjx.Model cannot access model elements by name using the same intuitive API as mujoco.MjModel. Currently, users must either:

  • Use mjx.name2id(model, mujoco.mjtObj.mjOBJ_BODY, "name") (verbose, 3 arguments)
  • Carry around both mjx.Model and mujoco.MjModel objects together
  • This is particularly problematic when using APIs like BraxDomainRandomizationVmapWrapper which only pass mjx.Model to the randomization function.

What this PR does

Adds named element access methods to mjx.Model that mirror the mujoco.MjModel API:

# Before (verbose)
body_id = mjx.name2id(model, mujoco.mjtObj.mjOBJ_BODY, "body_name")

# After (clean, consistent with MjModel)
body_id = model.body("body_name").id

Added methods:

body(), joint() (alias: jnt()), geom(), site()
camera() (alias: cam()), mesh(), actuator(), sensor()
tendon(), equality() (alias: eq()), numeric(), tuple(), key() (alias: keyframe())

Implementation:

  • Single _get_element() helper method handles all lookups (DRY)
  • Uses mujoco.mjtObj enum for type-safety
  • Returns lightweight view dataclasses (BodyView, JointView, etc.) with id and name attributes
  • Same name-parsing logic as existing support.name2id()

Testing :

import mujoco
from mujoco import mjx

model = mujoco.MjModel.from_xml_string("""
    <mujoco>
      <worldbody>
        <body name="body"><joint name="joint"/><geom name="geom" size="0.1"/></body>
      </worldbody>
      <actuator><motor name="motor" joint="joint"/></actuator>
    </mujoco>
""")
modelx = mjx.put_model(model)

assert model.body("body").id == modelx.body("body").id
assert model.joint("joint").id == modelx.joint("joint").id
assert model.actuator("motor").id == modelx.actuator("motor").id
assert model.geom("geom").id == modelx.geom("geom").id

@google-cla
Copy link

google-cla bot commented Dec 29, 2025

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

@Yuvraj-cyborg
Copy link
Author

Yuvraj-cyborg commented Dec 29, 2025

cc: @hartikainen

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

Named element accessors for mjx.Model

1 participant

@Yuvraj-cyborg