Skip to content
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

Python data structure inconsistencies and usability improvements #2402

Open
hartikainen opened this issue Feb 3, 2025 · 1 comment
Open

Python data structure inconsistencies and usability improvements #2402

hartikainen opened this issue Feb 3, 2025 · 1 comment
Assignees
@quagla
Labels
enhancement New feature or request

Comments

@hartikainen
Copy link
Contributor

hartikainen commented Feb 3, 2025 •
edited
Loading

The feature, motivation and pitch

Hey folks,

I've been using the new MjSpec's Python api quite a lot recently, and find the object- and name-based data structure accessors super useful. The interfaces between MjSpec, Mj{Model,Data}, and mjx.{Model,Data}, however, are quite inconsistent, which decreases the user experience quite a bit. I think making the different interfaces consistent with each other and polishing some of the rough edges could improve the usability significantly. For some examples of what I mean exactly, see the code snippet and the inlined comments below.

Alternatives

All these have workarounds.

Additional context

Examples 
from absl import app
import mujoco
from mujoco import mjx


def main(args):
    xml=r"""
<mujoco>
  <worldbody>
    <body name="main">
      <geom name="main" size="0.15 0.15 0.15" mass="1" type="box"/>
      <freejoint/>
      <body name="box">
        <joint name="box" type="hinge" range="-1 +1"/>
        <geom name="box" size="0.15 0.15 0.15" mass="1" type="box"/>
      </body>
      <body name="sphere">
        <joint name="sphere" type="hinge" range="-1 +1"/>
        <geom name="sphere" size="0.15 0.15 0.15" mass="1" type="box"/>
      </body>
    </body>
  </worldbody>
</mujoco>
"""

    spec = mujoco.MjSpec.from_string(xml)

    sphere_body = spec.find_body("main")
    box_body = spec.find_body("box")
    # Why can't I do `spec.find_geom("sphere")` as I can do with bodies? Their names are unique, after all.
    # Or even better: `spec.geom("sphere")` as can  be done with `Mj{Model,Data}`.
    sphere_geom = next(g for g in spec.geoms if g.name == "sphere")
    box_geom = next(g for g in spec.geoms if g.name == "box")

    # Again, why can't I do `spec.joint("box")`? This feels inconsistent with `Mj{Model,Data}`.
    joint_box = next(j for j in spec.joints if j.name == "box")
    joint_sphere = next(j for j in spec.joints if j.name == "sphere")
    joints = [joint_box, joint_sphere]

    mj_model = spec.compile()
    mj_data = mujoco.MjData(mj_model)
    mjx_model = mjx.put_model(mj_model)
    mjx_data = mjx.put_data(mj_model, mj_data)

    mj_data.bind(joint_sphere).qpos  # This works.
    mj_data.bind(joint_sphere).qvel  # This too.
    mj_data.bind(joints).qpos  # Doesn't work with multiple joints though.
    mjx_data.bind(mjx_model, joint_sphere).qpos  # This doesn't, but I believe it's fixed in https://github.com/google-deepmind/mujoco/commit/3dcfad3293522f564eeeb68bc178edb5ffe9c823
    mjx_data.bind(mjx_model, joint_sphere).qvel  # This doesn't either. I don't think it's fixed like `qpos`.
    mjx_data.bind(mjx_model, joints).qpos  # Neither does this.

    mj_model.bind(joint_box).qposadr  # This works.
    mj_model.bind(joints).qposadr  # This does not: "TypeError: bind(): incompatible function arguments."
    mjx_model.bind(joint_box).qposadr  # Neither does this: "IndexError: too many indices for array: array is 1-dimensional, but 2 were indexed"
    mjx_model.bind(joints).qposadr  # Same here: "IndexError: too many indices for array: array is 1-dimensional, but 2 were indexed"

    # None of these work. Why not? Why can I access `MjData` fields with `MjSpec`
    # objects but not `MjModel` objects?
    mj_data.bind(mj_model.joint("sphere")).qpos
    mj_data.bind([mj_model.joint("sphere"), mj_model.joint("box")]).qpos

    # This does not work. Why not? Why can I do `MjModel.{joint,geom,body}` but not
    # `mjx.Model.{joint,geom,body}`?
    mjx_data.bind(mjx_model.joint("sphere")).qpos
    mjx_data.bind([mjx_model.joint("sphere"), mjx_model.joint("box")]).qpos

    # Neither of these works. They might be a nice user-friendly feature to have. Not
    # sure though.
    mj_model.joint(["sphere", "box"])
    mjx_model.joint(["sphere", "box"])


if __name__ == "__main__":
    app.run(main)
@hartikainen hartikainen added the enhancement New feature or request label Feb 3, 2025
@yuvaltassa
Copy link
Collaborator

Thanks @hartikainen, this is super useful!

@hartikainen hartikainen changed the title Python data structure inconsistensies and usability improvements Python data structure inconsistencies and usability improvements Feb 4, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Assignees

@quagla quagla

Labels
enhancement New feature or request
Projects
None yet
Milestone
No milestone
Development

No branches or pull requests
3 participants
@hartikainen @yuvaltassa @quagla