MjSpec.to_xml changes joint/state-vector ordering after attach round-trip

[robin-moss]

Intro

Hi!

I am a simulation engineer at Ocado, and we use MuJoCo for simulation, validation, and testing of robot control.

I am leaving Ocado next week, so @Stonary-HZ is the ongoing team contact for this issue after 24/07/26. They have access to this reproduction and can answer follow-up questions or validate proposed fixes.

My setup

OS: Linux, Ubuntu 24.04 (AMD64)
MuJoCo: 3.8.1, with the Python API (plus some C++ and custom pybind)

What's happening? What did you expect?

MjSpec.to_xml() appears to change the compiled model's joint order after exporting and re-importing a modular MJCF that uses <asset><model> and <attach>.

The source model attaches a static model first, then attaches a moving model under a frame. The moving model contains:

1. a slider body with a slide joint
2. a free body with a free joint

On first load/compile, the slide joint appears before the free joint in the compiled model's joint layout.

After:

exported_xml = source_spec.to_xml()
# Write `exported_xml` to file 
exported_spec = mujoco.MjSpec.from_file("exported.xml")
exported_model = exported_spec.compile()

The exported model has the free joint before the slide joint.

This was spotted when reloading a model and giving it a saved state, the freejoints ended up in unexpected locations.

In the reproduction, I set moving_slide to 0.5, save time + qpos + qvel from the original model, then apply that state to the exported/re-imported model.

Expected behaviour:

The joint ordering when importing the original and exported models to end up the same (even if the XML ordering is slightly different)

The MuJoCo docs for saving models say the XML writer attempts to generate the smallest MJCF file "guaranteed to compile into the same model", modulo negligible numeric differences. I believe the changing joint order violates that expectation.

Docs:

• https://mujoco.readthedocs.io/en/stable/modeling.html#saving-models
• https://mujoco.readthedocs.io/en/stable/XMLreference.html#body-joint

Steps for reproduction

1. Save the three MJCF xml files (same names as below)

• scene.xml
• static.xml
• moving.xml

2. Save the python script reproduce.py
3. Run reproduce.py

I've tried this script against 3.10.0 in an separate Python env and got the same results.

Script Output

MuJoCo version: 3.10.0
Source MJCF:    /home/robinmoss/tmp/joint_order/scene.xml
Exported MJCF:  /home/robinmoss/tmp/joint_order/exported.xml
State fields:   time | qpos | qvel
State size:     original=16, exported=16

Joint layout, in compiled joint-id order
model     id  body           joint              type   qpos  qvel/dof
--------  --  -------------  -----------------  -----  ----  --------
original  0   moving_slider  moving_slide       slide  0:1   0:1     
original  1   moving_free    moving_free_joint  free   1:8   1:7     
exported  0   moving_free    moving_free_joint  free   0:7   0:6     
exported  1   moving_slider  moving_slide       slide  7:8   6:7     

Body positions after applying the original state to the exported model
body           original xpos   restored xpos   max abs error
-------------  --------------  --------------  -------------
moving_slider  [ 0.5,  2,  0]  [ 0,  2,  0]    0.5          
moving_free    [ 1,  2,  0]    [ 0.5,  1,  2]  2            
Traceback (most recent call last):
  File "/home/robinmoss/tmp/joint_order/reproduce.py", line 310, in <module>
    main()
  File "/home/robinmoss/tmp/joint_order/reproduce.py", line 304, in main
    raise AssertionError(f"Round-trip invariants failed:\n\n{formatted_failures}")
AssertionError: Round-trip invariants failed:

1. Joint layout changed after MjSpec.to_xml() export/import.
Expected original layout:
  0: body='moving_slider', joint='moving_slide', type='slide', qpos=0:1, qvel/dof=0:1
  1: body='moving_free', joint='moving_free_joint', type='free', qpos=1:8, qvel/dof=1:7
Actual exported layout:
  0: body='moving_free', joint='moving_free_joint', type='free', qpos=0:7, qvel/dof=0:6
  1: body='moving_slider', joint='moving_slide', type='slide', qpos=7:8, qvel/dof=6:7

2. Body 'moving_slider' changed position after restoring the saved state.
Expected: [ 0.5,  2,  0]
Actual:   [ 0,  2,  0]
Max absolute error: 0.5

3. Body 'moving_free' changed position after restoring the saved state.
Expected: [ 1,  2,  0]
Actual:   [ 0.5,  1,  2]
Max absolute error: 2

Minimal model for reproduction

If you encountered the issue in a complex model, please simplify it as much as possible (while still reproducing the issue).

scene.xml

<mujoco model="export_state_order_reproduction">
  <asset>
    <model name="static_model" file="static.xml"/>
    <model name="moving_model" file="moving.xml"/>
  </asset>

  <worldbody>
    <!-- This direct attachment is required to reproduce the serialization order change. -->
    <attach model="static_model" prefix="static_"/>

    <frame pos="0 2 0">
      <attach model="moving_model" prefix="moving_"/>
    </frame>
  </worldbody>
</mujoco>

static.xml

<mujoco model="static_model">
  <worldbody>
    <body name="body">
      <geom type="box" size="0.1 0.1 0.1"/>
    </body>
  </worldbody>
</mujoco>

moving.xml

<mujoco model="moving_model">
  <worldbody>
    <body name="slider">
      <joint name="slide" type="slide" axis="1 0 0"/>
      <geom type="box" size="0.1 0.1 0.1" mass="1"/>
    </body>

    <frame pos="1 0 0">
      <body name="free">
        <freejoint name="free_joint"/>
        <geom type="box" size="0.1 0.1 0.1" mass="1"/>
      </body>
    </frame>
  </worldbody>
</mujoco>

Code required for reproduction

The script does:

1. Import the scene.xml (which attaches static and moving)
2. Updates the slide joints qpos (for testing state recovery)
3. Uses mj_getState to save current state
4. Captures current body position
5. Captures joint ordering
6. Exports the MJCF

from dataclasses import dataclass
from pathlib import Path

import mujoco
import numpy as np

ISSUE_DIR = Path(__file__).resolve().parent
SOURCE_MJCF = ISSUE_DIR / "scene.xml"
EXPORTED_MJCF = ISSUE_DIR / "exported.xml"

STATE_SIGNATURE = int(
    mujoco.mjtState.mjSTATE_TIME
    | mujoco.mjtState.mjSTATE_QPOS
    | mujoco.mjtState.mjSTATE_QVEL
)

RTOL = 0.0
ATOL = 1e-12

JOINT_TYPE_NAMES = {
    int(mujoco.mjtJoint.mjJNT_FREE): "free",
    int(mujoco.mjtJoint.mjJNT_BALL): "ball",
    int(mujoco.mjtJoint.mjJNT_SLIDE): "slide",
    int(mujoco.mjtJoint.mjJNT_HINGE): "hinge",
}

JOINT_WIDTHS = {
    int(mujoco.mjtJoint.mjJNT_FREE): (7, 6),
    int(mujoco.mjtJoint.mjJNT_BALL): (4, 3),
    int(mujoco.mjtJoint.mjJNT_SLIDE): (1, 1),
    int(mujoco.mjtJoint.mjJNT_HINGE): (1, 1),
}


@dataclass(frozen=True)
class JointLayoutRow:
    """One compiled joint-layout row."""

    index: int
    body: str
    joint: str
    joint_type: str
    qpos_address: int
    qpos_width: int
    dof_address: int
    dof_width: int

    @property
    def qpos_range(self) -> str:
        """Return the half-open qpos range used by this joint."""
        return f"{self.qpos_address}:{self.qpos_address + self.qpos_width}"

    @property
    def dof_range(self) -> str:
        """Return the half-open qvel/dof range used by this joint."""
        return f"{self.dof_address}:{self.dof_address + self.dof_width}"

    def signature(self) -> tuple[str, str, str, int, int, int, int]:
        """Return the fields that should survive XML export/import unchanged."""
        return (
            self.body,
            self.joint,
            self.joint_type,
            self.qpos_address,
            self.qpos_width,
            self.dof_address,
            self.dof_width,
        )


def object_name(model: mujoco.MjModel, object_type: mujoco.mjtObj, object_id: int) -> str:
    """Return a MuJoCo object name, using a readable placeholder if unnamed."""
    name = mujoco.mj_id2name(model, object_type, object_id)
    return name if name is not None else "<unnamed>"


def require_id(model: mujoco.MjModel, object_type: mujoco.mjtObj, name: str) -> int:
    """Return a MuJoCo object id, raising a useful error if the object is missing."""
    object_id = mujoco.mj_name2id(model, object_type, name)
    if object_id < 0:
        object_type_name = getattr(object_type, "name", str(object_type))
        raise ValueError(f"Could not find {object_type_name} named {name!r}")
    return object_id


def joint_layout(model: mujoco.MjModel) -> list[JointLayoutRow]:
    """Return the compiled joint layout in joint-id order."""
    rows: list[JointLayoutRow] = []

    for joint_id in range(model.njnt):
        joint_type = int(model.jnt_type[joint_id])
        qpos_width, dof_width = JOINT_WIDTHS[joint_type]
        rows.append(
            JointLayoutRow(
                index=joint_id,
                body=object_name(
                    model,
                    mujoco.mjtObj.mjOBJ_BODY,
                    int(model.jnt_bodyid[joint_id]),
                ),
                joint=object_name(model, mujoco.mjtObj.mjOBJ_JOINT, joint_id),
                joint_type=JOINT_TYPE_NAMES[joint_type],
                qpos_address=int(model.jnt_qposadr[joint_id]),
                qpos_width=qpos_width,
                dof_address=int(model.jnt_dofadr[joint_id]),
                dof_width=dof_width,
            )
        )

    return rows


def body_position(model: mujoco.MjModel, data: mujoco.MjData, body_name: str) -> np.ndarray:
    """Return a copy of one body's world position."""
    body_id = require_id(model, mujoco.mjtObj.mjOBJ_BODY, body_name)
    return data.xpos[body_id].copy()


def format_vector(vector: np.ndarray) -> str:
    """Format a small vector compactly and consistently."""
    return "[" + ", ".join(f"{value: .6g}" for value in vector) + "]"


def format_table(title: str, headers: tuple[str, ...], rows: list[tuple[object, ...]]) -> str:
    """Return a simple aligned text table."""
    rendered_rows = [[str(cell) for cell in row] for row in rows]
    rendered_headers = list(headers)
    widths = [
        max(len(rendered_headers[column]), *(len(row[column]) for row in rendered_rows))
        for column in range(len(rendered_headers))
    ]

    def render_row(row: list[str]) -> str:
        return "  ".join(cell.ljust(widths[column]) for column, cell in enumerate(row))

    lines = [title, render_row(rendered_headers), render_row(["-" * width for width in widths])]
    lines.extend(render_row(row) for row in rendered_rows)
    return "\n".join(lines)


def joint_layout_table(original_layout: list[JointLayoutRow], exported_layout: list[JointLayoutRow]) -> str:
    """Return a formatted before/after joint-layout table."""
    rows: list[tuple[object, ...]] = []
    for label, layout in (("original", original_layout), ("exported", exported_layout)):
        rows.extend(
            (
                label,
                row.index,
                row.body,
                row.joint,
                row.joint_type,
                row.qpos_range,
                row.dof_range,
            )
            for row in layout
        )

    return format_table(
        "Joint layout, in compiled joint-id order",
        ("model", "id", "body", "joint", "type", "qpos", "qvel/dof"),
        rows,
    )


def body_position_table(
    original_positions: dict[str, np.ndarray],
    restored_positions: dict[str, np.ndarray],
) -> str:
    """Return a formatted original/restored body-position table."""
    rows = []
    for body_name, original_position in original_positions.items():
        restored_position = restored_positions[body_name]
        max_abs_error = float(np.max(np.abs(restored_position - original_position)))
        rows.append(
            (
                body_name,
                format_vector(original_position),
                format_vector(restored_position),
                f"{max_abs_error:.6g}",
            )
        )

    return format_table(
        "Body positions after applying the original state to the exported model",
        ("body", "original xpos", "restored xpos", "max abs error"),
        rows,
    )


def layout_signature(layout: list[JointLayoutRow]) -> str:
    """Return a compact multi-line signature for assertion messages."""
    return "\n".join(
        (
            f"  {row.index}: body={row.body!r}, joint={row.joint!r}, "
            f"type={row.joint_type!r}, qpos={row.qpos_range}, qvel/dof={row.dof_range}"
        )
        for row in layout
    )


def collect_failures(
    original_layout: list[JointLayoutRow],
    exported_layout: list[JointLayoutRow],
    original_positions: dict[str, np.ndarray],
    restored_positions: dict[str, np.ndarray],
) -> list[str]:
    """Return all round-trip invariant failures found in the reproduction."""
    failures: list[str] = []

    if [row.signature() for row in original_layout] != [row.signature() for row in exported_layout]:
        failures.append(
            "Joint layout changed after MjSpec.to_xml() export/import.\n"
            "Expected original layout:\n"
            f"{layout_signature(original_layout)}\n"
            "Actual exported layout:\n"
            f"{layout_signature(exported_layout)}"
        )

    for body_name, original_position in original_positions.items():
        restored_position = restored_positions[body_name]
        if not np.allclose(restored_position, original_position, rtol=RTOL, atol=ATOL):
            max_abs_error = float(np.max(np.abs(restored_position - original_position)))
            failures.append(
                f"Body {body_name!r} changed position after restoring the saved state.\n"
                f"Expected: {format_vector(original_position)}\n"
                f"Actual:   {format_vector(restored_position)}\n"
                f"Max absolute error: {max_abs_error:.6g}"
            )

    return failures


def main() -> None:
    """Export, recompile, and check that the saved state keeps the same meaning."""
    source_spec = mujoco.MjSpec.from_file(str(SOURCE_MJCF))
    source_model = source_spec.compile()
    source_data = mujoco.MjData(source_model)

    slide_joint_id = require_id(source_model, mujoco.mjtObj.mjOBJ_JOINT, "moving_slide")
    slide_qpos_address = int(source_model.jnt_qposadr[slide_joint_id])
    source_data.qpos[slide_qpos_address] = 0.5
    mujoco.mj_forward(source_model, source_data)

    saved_state = np.empty(
        mujoco.mj_stateSize(source_model, STATE_SIGNATURE),
        dtype=np.float64,
    )
    mujoco.mj_getState(source_model, source_data, saved_state, STATE_SIGNATURE)

    tracked_bodies = ("moving_slider", "moving_free")
    original_positions = {
        body_name: body_position(source_model, source_data, body_name)
        for body_name in tracked_bodies
    }
    original_layout = joint_layout(source_model)

    EXPORTED_MJCF.write_text(source_spec.to_xml(), encoding="utf-8")
    exported_spec = mujoco.MjSpec.from_file(str(EXPORTED_MJCF))
    exported_model = exported_spec.compile()
    exported_data = mujoco.MjData(exported_model)
    exported_state_size = mujoco.mj_stateSize(exported_model, STATE_SIGNATURE)

    print(f"MuJoCo version: {mujoco.__version__}")
    print(f"Source MJCF:    {SOURCE_MJCF}")
    print(f"Exported MJCF:  {EXPORTED_MJCF}")
    print(f"State fields:   time | qpos | qvel")
    print(f"State size:     original={saved_state.size}, exported={exported_state_size}")
    print()
    print(joint_layout_table(original_layout, joint_layout(exported_model)))

    if exported_state_size != saved_state.size:
        raise AssertionError(
            "Exported model has a different state size for time|qpos|qvel.\n"
            f"Original state size: {saved_state.size}\n"
            f"Exported state size: {exported_state_size}"
        )

    mujoco.mj_setState(exported_model, exported_data, saved_state, STATE_SIGNATURE)
    mujoco.mj_forward(exported_model, exported_data)

    exported_layout = joint_layout(exported_model)
    restored_positions = {
        body_name: body_position(exported_model, exported_data, body_name)
        for body_name in tracked_bodies
    }

    print()
    print(body_position_table(original_positions, restored_positions))

    failures = collect_failures(
        original_layout,
        exported_layout,
        original_positions,
        restored_positions,
    )
    if failures:
        formatted_failures = "\n\n".join(f"{index}. {failure}" for index, failure in enumerate(failures, 1))
        raise AssertionError(f"Round-trip invariants failed:\n\n{formatted_failures}")

    print("\nPASS: XML export/import preserved joint layout and saved-state semantics.")


if __name__ == "__main__":
    main()

Confirmations

• I searched the latest documentation thoroughly before posting.
• I searched previous Issues and Discussions, I am certain this has not been raised before.

[anishesg]

Took a look, #3402 should address it.

[robin-moss]

Thanks @anishesg, We'll try and validate it locally as well :)

[robin-moss]

We can confirm the linked PR does fix the issue, thanks