Repository: newton-physics/newton
Pull request #116: Expand randomizable model arrays in MuJoCo Solver, add initial support for runtime model updates
Author: login: adenzler-nvidia
Created: 2025-05-16T15:17:22Z | Merged: 2025-05-23T16:37:10Z

This is a very small first step towards getting model updates to work. It does:
- expand all the model fields that can be randomized right now
- fill body_com, body_mass, body_inertia with the actual values from the modelbuilder instead of the values from the first env.
- add initial support for runtime updates using the notify function.



Diff:
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 994730e0af..cc0473d50b 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -33,7 +33,7 @@ jobs:
         with:
           python-version-file: ".python-version"
       - name: Run Tests
-        run: uv run --all-extras -m newton.tests --junit-report-xml rspec.xml --coverage --coverage-xml coverage.xml
+        run: uv run --all-extras -m newton.tests --junit-report-xml rspec.xml --coverage --coverage-xml coverage.xml --serial-fallback
       - name: Test Summary
         uses: test-summary/action@31493c76ec9e7aa675f1585d3ed6f1da69269a86
         with:
diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml
index 9a7f54bd5c..8a960f0dc0 100644
--- a/.gitlab-ci.yml
+++ b/.gitlab-ci.yml
@@ -101,7 +101,7 @@ linting:
     - when: manual # If not auto-triggered, allow any pipeline to run this job manually
       allow_failure: true
   script:
-    - uv run --all-extras -m newton.tests --junit-report-xml rspec.xml --coverage --coverage-xml coverage.xml --failfast
+    - uv run --all-extras -m newton.tests --junit-report-xml rspec.xml --coverage --coverage-xml coverage.xml --failfast --serial-fallback
 
 linux-x86_64 test:
   image: ghcr.io/astral-sh/uv:bookworm
diff --git a/newton/core/types.py b/newton/core/types.py
index 438ae7d197..33c0b6aca9 100644
--- a/newton/core/types.py
+++ b/newton/core/types.py
@@ -347,6 +347,18 @@ def __hash__(self) -> int:
         return hash((tuple(np.array(self.vertices).flatten()), tuple(np.array(self.indices).flatten()), self.is_solid))
 
 
+# model update flags - used for solver.notify_model_update()
+NOTIFY_FLAG_JOINT_PROPERTIES = wp.constant(1 << 0)  # joint_q, joint_X_p, joint_X_c
+NOTIFY_FLAG_JOINT_AXIS_PROPERTIES = wp.constant(
+    1 << 1
+)  # joint_target, joint_target_ke, joint_target_kd, joint_axis_mode, joint_limit_upper, joint_limit_lower, joint_limit_ke, joint_limit_kd
+NOTIFY_FLAG_DOF_PROPERTIES = wp.constant(1 << 2)  # joint_qd, joint_f, joint_armature
+NOTIFY_FLAG_BODY_PROPERTIES = wp.constant(1 << 3)  # body_q, body_qd
+NOTIFY_FLAG_BODY_INERTIAL_PROPERTIES = wp.constant(
+    1 << 4
+)  # body_com, body_inertia, body_inv_inertia, body_mass, body_inv_mass
+NOTIFY_FLAG_SHAPE_PROPERTIES = wp.constant(1 << 5)  # shape_transform, shape_geo
+
 __all__ = [
     "GEO_BOX",
     "GEO_CAPSULE",
@@ -369,6 +381,11 @@ def __hash__(self) -> int:
     "JOINT_PRISMATIC",
     "JOINT_REVOLUTE",
     "JOINT_UNIVERSAL",
+    "NOTIFY_FLAG_BODY_INERTIAL_PROPERTIES",
+    "NOTIFY_FLAG_BODY_PROPERTIES",
+    "NOTIFY_FLAG_DOF_PROPERTIES",
+    "NOTIFY_FLAG_JOINT_PROPERTIES",
+    "NOTIFY_FLAG_SHAPE_PROPERTIES",
     "PARTICLE_FLAG_ACTIVE",
     "SDF",
     "SHAPE_FLAG_COLLIDE_GROUND",
diff --git a/newton/solvers/solver.py b/newton/solvers/solver.py
index 1a218ed701..d13418b70d 100644
--- a/newton/solvers/solver.py
+++ b/newton/solvers/solver.py
@@ -269,3 +269,12 @@ def step(self, model: Model, state_in: State, state_out: State, control: Control
             dt (float): The time step (typically in seconds).
         """
         raise NotImplementedError()
+
+    def notify_model_changed(self, flags: int):
+        """
+        Notify the solver that the model has changed.
+
+        Args:
+            flags (int): The flags that have changed.
+        """
+        pass
diff --git a/newton/solvers/solver_mujoco.py b/newton/solvers/solver_mujoco.py
index f0045be8a9..4321545cb7 100644
--- a/newton/solvers/solver_mujoco.py
+++ b/newton/solvers/solver_mujoco.py
@@ -15,14 +15,14 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 import warp as wp
 from typing_extensions import override
 
 import newton
-from newton.core import Contact, Control, Model, State
+from newton.core import Contact, Control, Model, State, types
 
 from .solver import SolverBase
 
@@ -498,6 +498,95 @@ def convert_body_xforms_to_warp_kernel(
     body_q[wbi] = wp.transform(pos, quat)
 
 
+@wp.kernel
+def update_body_mass_ipos_kernel(
+    body_com: wp.array(dtype=wp.vec3f),
+    body_mass: wp.array(dtype=float),
+    bodies_per_env: int,
+    up_axis: int,
+    body_mapping: wp.array(dtype=int),
+    # outputs
+    body_ipos: wp.array2d(dtype=wp.vec3f),
+    body_mass_out: wp.array2d(dtype=float),
+):
+    tid = wp.tid()
+    worldid = wp.tid() // bodies_per_env
+    index_in_env = wp.tid() % bodies_per_env
+    mjc_idx = body_mapping[index_in_env]
+    if mjc_idx == -1:
+        return
+
+    # Update COM position
+    if up_axis == 1:
+        body_ipos[worldid, mjc_idx] = wp.vec3f(body_com[tid][0], -body_com[tid][2], body_com[tid][1])
+    else:
+        body_ipos[worldid, mjc_idx] = body_com[tid]
+
+    # Update mass
+    body_mass_out[worldid, mjc_idx] = body_mass[tid]
+
+
+@wp.kernel
+def update_body_inertia_kernel(
+    body_inertia: wp.array(dtype=wp.mat33f),
+    body_quat: wp.array2d(dtype=wp.quatf),
+    bodies_per_env: int,
+    body_mapping: wp.array(dtype=int),
+    # outputs
+    body_inertia_out: wp.array2d(dtype=wp.vec3f),
+    body_iquat_out: wp.array2d(dtype=wp.quatf),
+):
+    tid = wp.tid()
+    worldid = wp.tid() // bodies_per_env
+    index_in_env = wp.tid() % bodies_per_env
+    mjc_idx = body_mapping[index_in_env]
+    if mjc_idx == -1:
+        return
+
+    # Get inertia tensor and body orientation
+    I = body_inertia[tid]
+    # body_q = body_quat[worldid, mjc_idx]
+
+    # Calculate eigenvalues and eigenvectors
+    eigenvectors, eigenvalues = wp.eig3(I)
+
+    # Bubble sort for 3 elements in descending order
+    for i in range(2):
+        for j in range(2 - i):
+            if eigenvalues[j] < eigenvalues[j + 1]:
+                # Swap eigenvalues
+                temp_val = eigenvalues[j]
+                eigenvalues[j] = eigenvalues[j + 1]
+                eigenvalues[j + 1] = temp_val
+                # Swap eigenvectors
+                temp_vec = eigenvectors[j]
+                eigenvectors[j] = eigenvectors[j + 1]
+                eigenvectors[j + 1] = temp_vec
+
+    # this does not work yet, I think we are reporting in the wrong reference frame
+    # Convert eigenvectors to quaternion (xyzw format for mujoco)
+    # q = wp.quat_from_matrix(wp.mat33f(eigenvectors[0], eigenvectors[1], eigenvectors[2]))
+    # q = wp.normalize(q)
+
+    # Convert from wxyz to xyzw format and compose with body orientation
+    # q = wp.quat(q[1], q[2], q[3], q[0])
+
+    # Store results
+    body_inertia_out[worldid, mjc_idx] = eigenvalues
+    # body_iquat_out[worldid, mjc_idx] = q
+
+
+@wp.kernel
+def repeat_array_kernel(
+    src: wp.array(dtype=Any),
+    nelems_per_world: int,
+    dst: wp.array(dtype=Any),
+):
+    tid = wp.tid()
+    src_idx = tid % nelems_per_world
+    dst[tid] = src[src_idx]
+
+
 class MuJoCoSolver(SolverBase):
     """
     This solver provides an interface to simulate physics using the `MuJoCo <https://github.com/google-deepmind/mujoco>`_ physics engine,
@@ -576,7 +665,7 @@ def __init__(
             self.use_mujoco = use_mujoco
             if separate_envs_to_worlds is None:
                 separate_envs_to_worlds = not use_mujoco
-            (self.mjw_model, self.mjw_data, self.mj_model, self.mj_data) = self.convert_to_mjc(
+            self.convert_to_mjc(
                 model,
                 disableflags=disableflags,
                 separate_envs_to_worlds=separate_envs_to_worlds,
@@ -625,6 +714,11 @@ def step(self, model: Model, state_in: State, state_out: State, control: Control
         self._step += 1
         return state_out
 
+    @override
+    def notify_model_changed(self, flags: int):
+        if flags & types.NOTIFY_FLAG_BODY_INERTIAL_PROPERTIES:
+            self.update_model_inertial_properties()
+
     @staticmethod
     def _data_is_mjwarp(data):
         # Check if the data is a mujoco_warp Data object
@@ -796,8 +890,8 @@ def update_newton_state(model: Model, state: State, mj_data: MjWarpData | MjData
                 outputs=[state.body_q],
             )
 
-    @staticmethod
     def convert_to_mjc(
+        self,
         model: Model,
         state: State | None = None,
         *,
@@ -901,7 +995,7 @@ def pos2mjc(p):
         defaults.geom.solimp = geom_solimp
         defaults.geom.friction = geom_friction
         # defaults.geom.contype = 0
-        spec.compiler.inertiafromgeom = True
+        spec.compiler.inertiafromgeom = mujoco.mjtInertiaFromGeom.mjINERTIAFROMGEOM_AUTO
 
         if add_axes:
             # TODO figure out how to create noncolliding geoms
@@ -1154,6 +1248,7 @@ def add_geoms(warp_body_id: int, perm_position=False, incoming_xform=None):
                 mass=body_mass[child],
                 ipos=pos2mjc(body_com[child, :]),
                 fullinertia=[inertia[0, 0], inertia[1, 1], inertia[2, 2], inertia[0, 1], inertia[0, 2], inertia[1, 2]],
+                explicitinertial=True,
             )
             mj_bodies.append(body)
 
@@ -1281,7 +1376,7 @@ def add_geoms(warp_body_id: int, perm_position=False, incoming_xform=None):
 
             add_geoms(child, incoming_xform=child_tf)
 
-        m = spec.compile()
+        self.mj_model = spec.compile()
 
         if target_filename:
             import os
@@ -1290,36 +1385,179 @@ def add_geoms(warp_body_id: int, perm_position=False, incoming_xform=None):
                 f.write(spec.to_xml())
                 print(f"Saved mujoco model to {os.path.abspath(target_filename)}")
 
-        d = mujoco.MjData(m)
+        self.mj_data = mujoco.MjData(self.mj_model)
 
-        d.nefc = nefc_per_env
+        self.mj_data.nefc = nefc_per_env
 
-        m.opt.tolerance = tolerance
-        m.opt.ls_tolerance = ls_tolerance
-        m.opt.cone = cone
-        m.opt.iterations = iterations
-        m.opt.ls_iterations = ls_iterations
-        m.opt.integrator = integrator
-        m.opt.solver = solver
+        self.mj_model.opt.tolerance = tolerance
+        self.mj_model.opt.ls_tolerance = ls_tolerance
+        self.mj_model.opt.cone = cone
+        self.mj_model.opt.iterations = iterations
+        self.mj_model.opt.ls_iterations = ls_iterations
+        self.mj_model.opt.integrator = integrator
+        self.mj_model.opt.solver = solver
         # m.opt.disableflags = disableflags
-        m.opt.impratio = impratio
-        m.opt.jacobian = mujoco.mjtJacobian.mjJAC_AUTO
+        self.mj_model.opt.impratio = impratio
+        self.mj_model.opt.jacobian = mujoco.mjtJacobian.mjJAC_AUTO
 
-        MuJoCoSolver.update_mjc_data(d, model, state)
+        MuJoCoSolver.update_mjc_data(self.mj_data, model, state)
 
-        mujoco.mj_forward(m, d)
+        # fill some MjWarp model fields that outdated after update_mjc_data.
+        # just setting qpos0 to d.qpos leads to weird behavior here, needs
+        # to be investigated.
+
+        mujoco.mj_forward(self.mj_model, self.mj_data)
 
         with wp.ScopedDevice(model.device):
             # add axis_to_actuator mapping to the Newton model
             model.axis_to_actuator = wp.array(axis_to_actuator, dtype=wp.int32)  # pyright: ignore[reportAttributeAccessIssue]
 
-            mj_model = mujoco_warp.put_model(m)
+            self.mjw_model = mujoco_warp.put_model(self.mj_model)
             if separate_envs_to_worlds:
                 nworld = model.num_envs
             else:
                 nworld = 1
-            nconmax = max(model.rigid_contact_max, d.ncon * nworld)  # this avoids error in mujoco.
-            njmax = max(nworld * nefc_per_env, nworld * d.nefc)
-            mj_data = mujoco_warp.put_data(m, d, nworld=nworld, nconmax=nconmax, njmax=njmax)
 
-        return mj_model, mj_data, m, d
+            # expand model fields that can be expanded:
+            self.expand_model_fields(self.mjw_model, nworld)
+
+            # complete the body mapping
+            size = max(body_mapping.keys()) + 1
+            arr = np.full(size, -1, dtype=int)
+            for k, v in body_mapping.items():
+                if k != -1:
+                    arr[k] = v
+
+            self.body_mapping = wp.array(arr, dtype=int)
+
+            # now fill with all the data from the Newton model.
+            flags = types.NOTIFY_FLAG_BODY_INERTIAL_PROPERTIES
+            self.notify_model_changed(flags)
+
+            # TODO find better heuristics to determine nconmax and njmax
+            nconmax = max(model.rigid_contact_max, self.mj_data.ncon * nworld)  # this avoids error in mujoco.
+            njmax = max(nworld * nefc_per_env, nworld * self.mj_data.nefc)
+            self.mjw_data = mujoco_warp.put_data(
+                self.mj_model, self.mj_data, nworld=nworld, nconmax=nconmax, njmax=njmax
+            )
+
+    def expand_model_fields(self, mj_model: MjWarpModel, nworld: int):
+        if nworld == 1:
+            return
+
+        model_fields_to_expand = [
+            # "qpos0",
+            # "qpos_spring",
+            # "body_pos",
+            # "body_quat",
+            "body_ipos",
+            # "body_iquat",
+            "body_mass",
+            # "body_subtreemass",
+            # "subtree_mass",
+            "body_inertia",
+            # "body_invweight0",
+            # "body_gravcomp",
+            # "jnt_solref",
+            # "jnt_solimp",
+            # "jnt_pos",
+            # "jnt_axis",
+            # "jnt_stiffness",
+            # "jnt_range",
+            # "jnt_actfrcrange",
+            # "jnt_margin",
+            # "dof_armature",
+            # "dof_damping",
+            # "dof_invweight0",
+            # "dof_frictionloss",
+            # "dof_solimp",
+            # "dof_solref",
+            # "geom_matid",
+            # "geom_solmix",
+            # "geom_solref",
+            # "geom_solimp",
+            # "geom_size",
+            # "geom_rbound",
+            # "geom_pos",
+            # "geom_quat",
+            # "geom_friction",
+            # "geom_margin",
+            # "geom_gap",
+            # "geom_rgba",
+            # "site_pos",
+            # "site_quat",
+            # "cam_pos",
+            # "cam_quat",
+            # "cam_poscom0",
+            # "cam_pos0",
+            # "cam_mat0",
+            # "light_pos",
+            # "light_dir",
+            # "light_poscom0",
+            # "light_pos0",
+            # "eq_solref",
+            # "eq_solimp",
+            # "eq_data",
+            # "actuator_dynprm",
+            # "actuator_gainprm",
+            # "actuator_biasprm",
+            # "actuator_ctrlrange",
+            # "actuator_forcerange",
+            # "actuator_actrange",
+            # "actuator_gear",
+            # "pair_solref",
+            # "pair_solreffriction",
+            # "pair_solimp",
+            # "pair_margin",
+            # "pair_gap",
+            # "pair_friction",
+            # "tendon_solref_lim",
+            # "tendon_solimp_lim",
+            # "tendon_range",
+            # "tendon_margin",
+            # "tendon_length0",
+            # "tendon_invweight0",
+            # "mat_rgba",
+        ]
+
+        def tile(x):
+            # Create new array with same shape but first dim multiplied by nworld
+            new_shape = list(x.shape)
+            new_shape[0] = nworld
+            wp_array = {1: wp.array, 2: wp.array2d, 3: wp.array3d, 4: wp.array4d}[len(new_shape)]
+            dst = wp_array(shape=new_shape, dtype=x.dtype, device=x.device)
+
+            # Flatten arrays for kernel
+            src_flat = x.flatten()
+            dst_flat = dst.flatten()
+
+            # Launch kernel to repeat data - one thread per destination element
+            n_elems_per_world = dst_flat.shape[0] // nworld
+            wp.launch(
+                repeat_array_kernel, dim=dst_flat.shape[0], inputs=[src_flat, n_elems_per_world], outputs=[dst_flat]
+            )
+            return dst
+
+        for field in mj_model.__dataclass_fields__:
+            if field in model_fields_to_expand:
+                array = getattr(mj_model, field)
+                setattr(mj_model, field, tile(array))
+
+    def update_model_inertial_properties(self):
+        bodies_per_env = self.model.body_count // self.model.num_envs
+
+        wp.launch(
+            update_body_mass_ipos_kernel,
+            dim=self.model.body_count,
+            inputs=[self.model.body_com, self.model.body_mass, bodies_per_env, self.model.up_axis, self.body_mapping],
+            outputs=[self.mjw_model.body_ipos, self.mjw_model.body_mass],
+            device=self.model.device,
+        )
+
+        wp.launch(
+            update_body_inertia_kernel,
+            dim=self.model.body_count,
+            inputs=[self.model.body_inertia, self.mjw_model.body_quat, bodies_per_env, self.body_mapping],
+            outputs=[self.mjw_model.body_inertia, self.mjw_model.body_iquat],
+            device=self.model.device,
+        )
diff --git a/newton/tests/test_mujoco_solver.py b/newton/tests/test_mujoco_solver.py
new file mode 100644
index 0000000000..49770de1e2
--- /dev/null
+++ b/newton/tests/test_mujoco_solver.py
@@ -0,0 +1,442 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 The Newton Developers
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os  # For path manipulation
+import time  # Added for the interactive loop
+import unittest
+
+import numpy as np  # For numerical operations and random values
+import warp as wp
+
+import newton
+from newton.core import types
+from newton.solvers import MuJoCoSolver
+
+# Import the kernels for coordinate conversion
+from newton.utils import SimRendererOpenGL
+
+
+class TestMuJoCoSolver(unittest.TestCase):
+    def setUp(self):
+        """Set up a model with multiple environments, each with a free body and an articulated tree."""
+        self.seed = 123
+        self.rng = np.random.default_rng(self.seed)
+
+        num_envs = 2
+        self.debug_stage_path = "newton/tests/test_mujoco_render.usda"
+
+        template_builder = newton.ModelBuilder()
+        shape_cfg = newton.ModelBuilder.ShapeConfig(density=1000.0)  # Define ShapeConfig
+
+        # --- Free-floating body (e.g., a box) ---
+        free_body_initial_pos = wp.transform((0.5, 0.5, 0.0), wp.quat_identity())
+        free_body_idx = template_builder.add_body(mass=0.2)
+        template_builder.add_joint_free(child=free_body_idx, parent_xform=free_body_initial_pos)
+        template_builder.add_shape_box(
+            body=free_body_idx,
+            xform=wp.transform(),  # Shape at body's local origin
+            hx=0.1,
+            hy=0.1,
+            hz=0.1,
+            cfg=shape_cfg,
+        )
+
+        # --- Articulated tree (3 bodies) ---
+        link_radius = 0.05
+        link_half_length = 0.15
+        tree_root_initial_pos_y = link_half_length * 2.0
+        tree_root_initial_transform = wp.transform((0.0, tree_root_initial_pos_y, 0.0), wp.quat_identity())
+
+        body1_idx = template_builder.add_body(mass=0.1)
+        template_builder.add_joint_free(child=body1_idx, parent_xform=tree_root_initial_transform)
+        template_builder.add_shape_capsule(
+            body=body1_idx,
+            xform=wp.transform(),  # Shape at body's local origin
+            radius=link_radius,
+            half_height=link_half_length,
+            cfg=shape_cfg,
+        )
+
+        body2_idx = template_builder.add_body(mass=0.1)
+        template_builder.add_shape_capsule(
+            body=body2_idx,
+            xform=wp.transform(),  # Shape at body's local origin
+            radius=link_radius,
+            half_height=link_half_length,
+            cfg=shape_cfg,
+        )
+        template_builder.add_joint_revolute(
+            parent=body1_idx,
+            child=body2_idx,
+            parent_xform=wp.transform((0.0, link_half_length, 0.0), wp.quat_identity()),
+            child_xform=wp.transform((0.0, -link_half_length, 0.0), wp.quat_identity()),
+            axis=(0.0, 0.0, 1.0),
+        )
+
+        body3_idx = template_builder.add_body(mass=0.1)
+        template_builder.add_shape_capsule(
+            body=body3_idx,
+            xform=wp.transform(),  # Shape at body's local origin
+            radius=link_radius,
+            half_height=link_half_length,
+            cfg=shape_cfg,
+        )
+        template_builder.add_joint_revolute(
+            parent=body2_idx,
+            child=body3_idx,
+            parent_xform=wp.transform((0.0, link_half_length, 0.0), wp.quat_identity()),
+            child_xform=wp.transform((0.0, -link_half_length, 0.0), wp.quat_identity()),
+            axis=(1.0, 0.0, 0.0),
+        )
+
+        self.builder = newton.ModelBuilder()
+        self.builder.add_shape_plane()
+
+        for i in range(num_envs):
+            env_transform = wp.transform((i * 2.0, 0.0, 0.0), wp.quat_identity())
+            self.builder.add_builder(
+                template_builder, xform=env_transform, update_num_env_count=True, separate_collision_group=True
+            )
+
+        try:
+            if self.builder.num_envs == 0 and num_envs > 0:
+                self.builder.num_envs = num_envs
+            self.model = self.builder.finalize()
+            if self.model.num_envs != num_envs:
+                print(f"Warning: Model.num_envs ({self.model.num_envs}) does not match expected num_envs ({num_envs}).")
+        except Exception as e:
+            self.fail(f"Model finalization failed: {e}")
+
+        self.state_in = self.model.state()
+        self.state_out = self.model.state()
+        self.control = self.model.control()
+        self.contacts = self.model.contact()
+
+    def _run_substeps_for_frame(self, sim_dt, sim_substeps):
+        """Helper method to run simulation substeps for one rendered frame."""
+        for _ in range(sim_substeps):
+            self.solver.step(self.model, self.state_in, self.state_out, self.control, self.contacts, sim_dt)
+            self.state_in, self.state_out = self.state_out, self.state_in  # Output becomes input for next substep
+
+    def test_setup_completes(self):
+        """
+        Tests if the setUp method completes successfully.
+        This implicitly tests model creation, finalization, solver, and renderer initialization.
+        """
+        self.assertTrue(True, "setUp method completed.")
+
+    def test_randomize_body_mass(self):
+        """
+        Tests if the body mass is randomized correctly and updates properly after simulation steps.
+        """
+        # Randomize masses for all bodies in all environments
+        new_masses = self.rng.uniform(1.0, 10.0, size=self.model.body_count)
+        self.model.body_mass.assign(new_masses)
+
+        # Initialize solver
+        solver = MuJoCoSolver(self.model, ls_iterations=1, iterations=1, disable_contacts=True)
+
+        # Check that masses were transferred correctly
+        bodies_per_env = self.model.body_count // self.model.num_envs
+        body_mapping = solver.body_mapping.numpy()
+        for env_idx in range(self.model.num_envs):
+            for body_idx in range(bodies_per_env):
+                newton_idx = env_idx * bodies_per_env + body_idx
+                mjc_idx = body_mapping[body_idx]
+                if mjc_idx != -1:  # Skip unmapped bodies
+                    self.assertAlmostEqual(
+                        new_masses[newton_idx],
+                        solver.mjw_model.body_mass.numpy()[env_idx, mjc_idx],
+                        places=6,
+                        msg=f"Mass mismatch for body {body_idx} in environment {env_idx}",
+                    )
+
+        # Run a simulation step
+        solver.step(self.model, self.state_in, self.state_out, self.control, self.contacts, 0.01)
+        self.state_in, self.state_out = self.state_out, self.state_in
+
+        # Update masses again
+        updated_masses = self.rng.uniform(1.0, 10.0, size=self.model.body_count)
+        self.model.body_mass.assign(updated_masses)
+
+        # Notify solver of mass changes
+        solver.notify_model_changed(types.NOTIFY_FLAG_BODY_INERTIAL_PROPERTIES)
+
+        # Check that updated masses were transferred correctly
+        for env_idx in range(self.model.num_envs):
+            for body_idx in range(bodies_per_env):
+                newton_idx = env_idx * bodies_per_env + body_idx
+                mjc_idx = body_mapping[body_idx]
+                if mjc_idx != -1:  # Skip unmapped bodies
+                    self.assertAlmostEqual(
+                        updated_masses[newton_idx],
+                        solver.mjw_model.body_mass.numpy()[env_idx, mjc_idx],
+                        places=6,
+                        msg=f"Updated mass mismatch for body {body_idx} in environment {env_idx}",
+                    )
+
+    def test_randomize_body_com(self):
+        """
+        Tests if the body center of mass is randomized correctly and updates properly after simulation steps.
+        """
+        # Randomize COM for all bodies in all environments
+        new_coms = self.rng.uniform(-1.0, 1.0, size=(self.model.body_count, 3))
+        self.model.body_com.assign(new_coms)
+
+        # Initialize solver
+        solver = MuJoCoSolver(self.model, ls_iterations=1, iterations=1, disable_contacts=True, nefc_per_env=1)
+
+        # Check that COM positions were transferred correctly
+        bodies_per_env = self.model.body_count // self.model.num_envs
+        body_mapping = solver.body_mapping.numpy()
+        for env_idx in range(self.model.num_envs):
+            for body_idx in range(bodies_per_env):
+                newton_idx = env_idx * bodies_per_env + body_idx
+                mjc_idx = body_mapping[body_idx]
+                if mjc_idx != -1:  # Skip unmapped bodies
+                    newton_pos = new_coms[newton_idx]
+                    mjc_pos = solver.mjw_model.body_ipos.numpy()[env_idx, mjc_idx]
+
+                    # Convert positions based on up_axis
+                    if self.model.up_axis == 1:  # Y-axis up
+                        expected_pos = np.array([newton_pos[0], -newton_pos[2], newton_pos[1]])
+                    else:  # Z-axis up
+                        expected_pos = newton_pos
+
+                    for dim in range(3):
+                        self.assertAlmostEqual(
+                            expected_pos[dim],
+                            mjc_pos[dim],
+                            places=6,
+                            msg=f"COM position mismatch for body {body_idx} in environment {env_idx}, dimension {dim}",
+                        )
+
+        # Run a simulation step
+        solver.step(self.model, self.state_in, self.state_out, self.control, self.contacts, 0.01)
+        self.state_in, self.state_out = self.state_out, self.state_in
+
+        # Update COM positions again
+        updated_coms = self.rng.uniform(-1.0, 1.0, size=(self.model.body_count, 3))
+        self.model.body_com.assign(updated_coms)
+
+        # Notify solver of COM changes
+        solver.notify_model_changed(types.NOTIFY_FLAG_BODY_INERTIAL_PROPERTIES)
+
+        # Check that updated COM positions were transferred correctly
+        for env_idx in range(self.model.num_envs):
+            for body_idx in range(bodies_per_env):
+                newton_idx = env_idx * bodies_per_env + body_idx
+                mjc_idx = body_mapping[body_idx]
+                if mjc_idx != -1:  # Skip unmapped bodies
+                    newton_pos = updated_coms[newton_idx]
+                    mjc_pos = solver.mjw_model.body_ipos.numpy()[env_idx, mjc_idx]
+
+                    # Convert positions based on up_axis
+                    if self.model.up_axis == 1:  # Y-axis up
+                        expected_pos = np.array([newton_pos[0], -newton_pos[2], newton_pos[1]])
+                    else:  # Z-axis up
+                        expected_pos = newton_pos
+
+                    for dim in range(3):
+                        self.assertAlmostEqual(
+                            expected_pos[dim],
+                            mjc_pos[dim],
+                            places=6,
+                            msg=f"Updated COM position mismatch for body {body_idx} in environment {env_idx}, dimension {dim}",
+                        )
+
+    def test_randomize_body_inertia(self):
+        """
+        Tests if the body inertia is randomized correctly.
+        """
+        # Randomize inertia tensors for all bodies in all environments
+        # Simple inertia tensors that satisfy triangle inequality
+
+        new_inertias = np.zeros((self.model.body_count, 3, 3))
+        bodies_per_env = self.model.body_count // self.model.num_envs
+        for i in range(self.model.body_count):
+            env_idx = i // bodies_per_env
+            if env_idx == 0:
+                # First environment: ensure a + b > c with random values
+                a = 2.0 + self.rng.uniform(0.0, 0.5)
+                b = 3.0 + self.rng.uniform(0.0, 0.5)
+                c = min(a + b - 0.1, 4.0)  # Ensure a + b > c
+                new_inertias[i] = np.diag([a, b, c])
+            else:
+                # Second environment: ensure a + b > c with random values
+                a = 3.0 + self.rng.uniform(0.0, 0.5)
+                b = 4.0 + self.rng.uniform(0.0, 0.5)
+                c = min(a + b - 0.1, 5.0)  # Ensure a + b > c
+                new_inertias[i] = np.diag([a, b, c])
+        self.model.body_inertia.assign(new_inertias)
+
+        # Initialize solver
+        solver = MuJoCoSolver(self.model, iterations=1, ls_iterations=1, disable_contacts=True)
+
+        # Get body mapping once outside the loop
+        body_mapping = solver.body_mapping.numpy()
+
+        def check_inertias(inertias_to_check, msg_prefix=""):
+            for env_idx in range(self.model.num_envs):
+                for body_idx in range(bodies_per_env):
+                    newton_idx = env_idx * bodies_per_env + body_idx
+                    mjc_idx = body_mapping[body_idx]
+                    if mjc_idx != -1:  # Skip unmapped bodies
+                        newton_inertia = inertias_to_check[newton_idx]
+                        mjc_inertia = solver.mjw_model.body_inertia.numpy()[env_idx, mjc_idx]
+
+                        # Get eigenvalues of both tensors
+                        newton_eigvals = np.linalg.eigvalsh(newton_inertia)
+                        mjc_eigvals = mjc_inertia  # Already in diagonal form
+
+                        # Sort eigenvalues in descending order
+                        newton_eigvals.sort()
+                        newton_eigvals = newton_eigvals[::-1]
+
+                        for dim in range(3):
+                            self.assertAlmostEqual(
+                                newton_eigvals[dim],
+                                mjc_eigvals[dim],
+                                places=6,
+                                msg=f"{msg_prefix}Inertia eigenvalue mismatch for body {body_idx} in environment {env_idx}, dimension {dim}",
+                            )
+
+        # Check initial inertia tensors
+        check_inertias(new_inertias, "Initial ")
+
+        # Run a simulation step
+        solver.step(self.model, self.state_in, self.state_out, self.control, self.contacts, 0.01)
+        self.state_in, self.state_out = self.state_out, self.state_in
+
+        # Update inertia tensors again with new random values
+        updated_inertias = np.zeros((self.model.body_count, 3, 3))
+        for i in range(self.model.body_count):
+            env_idx = i // bodies_per_env
+            if env_idx == 0:
+                a = 2.5 + self.rng.uniform(0.0, 0.5)
+                b = 3.5 + self.rng.uniform(0.0, 0.5)
+                c = min(a + b - 0.1, 4.5)
+                updated_inertias[i] = np.diag([a, b, c])
+            else:
+                a = 3.5 + self.rng.uniform(0.0, 0.5)
+                b = 4.5 + self.rng.uniform(0.0, 0.5)
+                c = min(a + b - 0.1, 5.5)
+                updated_inertias[i] = np.diag([a, b, c])
+        self.model.body_inertia.assign(updated_inertias)
+
+        # Notify solver of inertia changes
+        solver.notify_model_changed(types.NOTIFY_FLAG_BODY_INERTIAL_PROPERTIES)
+
+        # Check updated inertia tensors
+        check_inertias(updated_inertias, "Updated ")
+
+    @unittest.skip("Trajectory rendering for debugging")
+    def test_render_trajectory(self):
+        """Simulates and renders a trajectory if solver and renderer are available."""
+        print("\nDebug: Starting test_render_trajectory...")
+
+        solver = None
+        renderer = None
+        substep_graph = None
+        use_cuda_graph = wp.get_device().is_cuda
+
+        try:
+            print("Debug: Attempting to initialize MuJoCoSolver for trajectory test...")
+            solver = MuJoCoSolver(self.model, iterations=10, ls_iterations=10)
+            print("Debug: MuJoCoSolver initialized successfully for trajectory test.")
+        except ImportError as e:
+            self.skipTest(f"MuJoCo or deps not installed. Skipping trajectory rendering: {e}")
+            return
+        except Exception as e:
+            self.skipTest(f"Error initializing MuJoCoSolver for trajectory test: {e}")
+            return
+
+        if self.debug_stage_path:
+            try:
+                print(f"Debug: Attempting to initialize SimRendererOpenGL (stage: {self.debug_stage_path})...")
+                stage_dir = os.path.dirname(self.debug_stage_path)
+                if stage_dir and not os.path.exists(stage_dir):
+                    os.makedirs(stage_dir)
+                    print(f"Debug: Created directory for stage: {stage_dir}")
+                renderer = SimRendererOpenGL(
+                    path=self.debug_stage_path, model=self.model, scaling=1.0, show_joints=True
+                )
+                print("Debug: SimRendererOpenGL initialized successfully for trajectory test.")
+            except ImportError as e:
+                self.skipTest(f"SimRendererOpenGL dependencies not met. Skipping trajectory rendering: {e}")
+                return
+            except Exception as e:
+                self.skipTest(f"Error initializing SimRendererOpenGL for trajectory test: {e}")
+                return
+        else:
+            self.skipTest("No debug_stage_path set. Skipping trajectory rendering.")
+            return
+
+        num_frames = 200
+        sim_substeps = 2
+        frame_dt = 1.0 / 60.0
+        sim_dt = frame_dt / sim_substeps
+        sim_time = 0.0
+
+        # Override self.solver for _run_substeps_for_frame if it was defined in setUp
+        # However, since we moved initialization here, we pass it directly or use the local var.
+        # For simplicity, let _run_substeps_for_frame use self.solver, so we assign the local one to it.
+        self.solver = solver  # Make solver accessible to _run_substeps_for_frame via self
+
+        if use_cuda_graph:
+            print(
+                f"Debug: CUDA device detected. Attempting to capture {sim_substeps} substeps with dt={sim_dt:.4f} into a CUDA graph..."
+            )
+            try:
+                with wp.ScopedCapture() as capture:
+                    self._run_substeps_for_frame(sim_dt, sim_substeps)
+                substep_graph = capture.graph
+                print("Debug: CUDA graph captured successfully.")
+            except Exception as e:
+                print(f"Debug: CUDA graph capture failed: {e}. Falling back to regular execution.")
+                substep_graph = None
+        else:
+            print("Debug: Not using CUDA graph (non-CUDA device or flag disabled).")
+
+        print(f"Debug: Simulating and rendering {num_frames} frames ({sim_substeps} substeps/frame)...")
+        print("       Press Ctrl+C in the console to stop early.")
+
+        try:
+            for frame_num in range(num_frames):
+                if frame_num % 20 == 0:
+                    print(f"Debug: Frame {frame_num}/{num_frames}, Sim time: {sim_time:.2f}s")
+
+                renderer.begin_frame(sim_time)
+                renderer.render(self.state_in)
+                renderer.end_frame()
+
+                if use_cuda_graph and substep_graph:
+                    wp.capture_launch(substep_graph)
+                else:
+                    self._run_substeps_for_frame(sim_dt, sim_substeps)
+
+                sim_time += frame_dt
+                time.sleep(0.016)
+
+        except KeyboardInterrupt:
+            print("\nDebug: Trajectory rendering stopped by user.")
+        except Exception as e:
+            self.fail(f"Error during trajectory rendering: {e}")
+        finally:
+            print("Debug: test_render_trajectory finished.")
+
+
+if __name__ == "__main__":
+    unittest.main()


Reviews:
user: login: adenzler-nvidia
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-16T15:20:58Z
user: login: adenzler-nvidia
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-16T15:21:13Z
user: login: adenzler-nvidia
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-16T15:45:39Z
user: login: adenzler-nvidia
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-22T08:41:38Z
user: login: adenzler-nvidia
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-22T08:43:16Z
user: login: eric-heiden
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-22T17:44:01Z
user: login: eric-heiden
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-22T17:48:26Z
user: login: eric-heiden
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-22T17:49:57Z
user: login: eric-heiden
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-22T17:54:07Z
user: login: eric-heiden
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-22T17:58:28Z
user: login: adenzler-nvidia
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-22T18:26:17Z
user: login: adenzler-nvidia
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-22T18:29:01Z
user: login: adenzler-nvidia
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-22T18:30:29Z
user: login: eric-heiden
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-22T19:38:33Z
user: login: eric-heiden
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-22T19:39:38Z
user: login: eric-heiden
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-22T19:40:20Z
user: login: adenzler-nvidia
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-23T08:45:42Z
user: login: adenzler-nvidia
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-23T08:46:28Z
user: login: eric-heiden
state: APPROVED
author_association: MEMBER
submitted_at: 2025-05-23T16:35:12Z

Review comments:
diff_hunk: @@ -229,6 +229,63 @@ def convert_warp_coords_to_mj_kernel(
             qvel[worldid, qd_i + i] = joint_qd[wqd_i + i]
 
 
+@wp.kernel
+def convert_joint_q_qpos0(
+    joint_q: wp.array(dtype=wp.float32),
+    joints_per_env: int,
+    up_axis: int,
+    joint_type: wp.array(dtype=wp.int32),
+    joint_q_start: wp.array(dtype=wp.int32),
+    joint_axis_dim: wp.array(dtype=wp.int32, ndim=2),
+    # outputs
+    qpos0: wp.array2d(dtype=wp.float32),
+):
+    jntid = wp.tid()
+    worldid = jntid // joints_per_env
+    jnt_in_env = jntid % joints_per_env
+
+    type = joint_type[jntid]
+    q_i = joint_q_start[jnt_in_env]
+    wq_i = joint_q_start[joints_per_env * worldid + jnt_in_env]
+
+    if type == newton.JOINT_FREE:
+        # convert position components
+        if up_axis == 1:
+            qpos0[worldid, q_i + 0] = joint_q[wq_i + 0]
+            qpos0[worldid, q_i + 1] = -joint_q[wq_i + 2]
+            qpos0[worldid, q_i + 2] = joint_q[wq_i + 1]
+        else:
+            for i in range(3):
+                qpos0[worldid, q_i + i] = joint_q[wq_i + i]
+
+        rot = wp.quat(
+            joint_q[wq_i + 3],
+            joint_q[wq_i + 4],
+            joint_q[wq_i + 5],
+            joint_q[wq_i + 6],
+        )
+        if up_axis == 1:
+            rot_y2z = wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), wp.pi * 0.5)
+            rot = rot_y2z * rot
+        # change quaternion order from xyzw to wxyz
+        qpos0[worldid, q_i + 3] = rot[3]
+        qpos0[worldid, q_i + 4] = rot[0]
+        qpos0[worldid, q_i + 5] = rot[1]
+        qpos0[worldid, q_i + 6] = rot[2]
+
+    elif type == newton.JOINT_BALL:
+        # change quaternion order from xyzw to wxyz
+        qpos0[worldid, q_i + 0] = joint_q[wq_i + 1]
+        qpos0[worldid, q_i + 1] = joint_q[wq_i + 2]
+        qpos0[worldid, q_i + 2] = joint_q[wq_i + 3]
+        qpos0[worldid, q_i + 3] = joint_q[wq_i + 0]
+    else:
+        axis_count = joint_axis_dim[jntid, 0] + joint_axis_dim[jntid, 1]
+        for i in range(axis_count):
+            # convert position components
+            qpos0[worldid, q_i + i] = joint_q[wq_i + i]
path: newton/solvers/solver_mujoco.py
user: login: adenzler-nvidia
body: This one is quite similar to the state update kernel right now, but I would prefer to keep them separate because we're most likely going to add additional stuff in here when we add support for randomizing more fields.
created_at: 2025-05-16T15:20:58Z
author_association: MEMBER
diff_hunk: @@ -229,6 +229,63 @@ def convert_warp_coords_to_mj_kernel(
             qvel[worldid, qd_i + i] = joint_qd[wqd_i + i]
 
 
+@wp.kernel
+def convert_joint_q_qpos0(
+    joint_q: wp.array(dtype=wp.float32),
+    joints_per_env: int,
+    up_axis: int,
+    joint_type: wp.array(dtype=wp.int32),
+    joint_q_start: wp.array(dtype=wp.int32),
+    joint_axis_dim: wp.array(dtype=wp.int32, ndim=2),
+    # outputs
+    qpos0: wp.array2d(dtype=wp.float32),
+):
+    jntid = wp.tid()
+    worldid = jntid // joints_per_env
+    jnt_in_env = jntid % joints_per_env
+
+    type = joint_type[jntid]
+    q_i = joint_q_start[jnt_in_env]
+    wq_i = joint_q_start[joints_per_env * worldid + jnt_in_env]
+
+    if type == newton.JOINT_FREE:
+        # convert position components
+        if up_axis == 1:
+            qpos0[worldid, q_i + 0] = joint_q[wq_i + 0]
+            qpos0[worldid, q_i + 1] = -joint_q[wq_i + 2]
+            qpos0[worldid, q_i + 2] = joint_q[wq_i + 1]
+        else:
+            for i in range(3):
+                qpos0[worldid, q_i + i] = joint_q[wq_i + i]
+
+        rot = wp.quat(
+            joint_q[wq_i + 3],
+            joint_q[wq_i + 4],
+            joint_q[wq_i + 5],
+            joint_q[wq_i + 6],
+        )
+        if up_axis == 1:
+            rot_y2z = wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), wp.pi * 0.5)
+            rot = rot_y2z * rot
path: newton/solvers/solver_mujoco.py
user: login: adenzler-nvidia
body: Do we still need the up-axis handling?
created_at: 2025-05-16T15:21:13Z
author_association: MEMBER
diff_hunk: @@ -1222,13 +1279,29 @@ def add_geoms(warp_body_id: int, perm_position=False, incoming_xform=None):
 
         MuJoCoSolver.update_mjc_data(d, model, state)
 
+        # fill some MjWarp model fields that outdated of update_mjc_data.
+
+        # update qpos0 to match joint_q. Technically qpos0 is the
+        # reference positions which could differ from the initial position
+        # though.
+        m.qpos0 = d.qpos
+        # qpos_spring defaults to qpos0
+        m.qpos_spring = d.qpos
path: newton/solvers/solver_mujoco.py
user: login: adenzler-nvidia
body: This part currently breaks the humanoid example but I cannot figure out why - apparently the update_mjc_data call updates the qpos values to something that I cannot use as qpos0. I'm not sure whether it's the joint_q inputs that are no good or what's going on.
created_at: 2025-05-16T15:45:39Z
author_association: MEMBER
diff_hunk: @@ -347,6 +347,13 @@ def __hash__(self) -> int:
         return hash((tuple(np.array(self.vertices).flatten()), tuple(np.array(self.indices).flatten()), self.is_solid))
 
 
+# model update flags - WIP grouping and naming still being worked on.
+NOTIFY_FLAG_JOINT_PROPERTIES = wp.constant(1 << 0)
+NOTIFY_FLAG_DOF_PROPERTIES = wp.constant(1 << 1)
+NOTIFY_FLAG_BODY_PROPERTIES = wp.constant(1 << 2)
+NOTIFY_FLAG_BODY_INERTIAL_PROPERTIES = wp.constant(1 << 3)
+NOTIFY_FLAG_SHAPE_PROPERTIES = wp.constant(1 << 4)
path: newton/core/types.py
user: login: adenzler-nvidia
body: This is not supposed to be complete - I think we'll figure out how many/what we'll need as we go along and add support for more model updates and testing.

But looking for feedback anyway as don't know other what the other solvers might need.
created_at: 2025-05-22T08:41:37Z
author_association: MEMBER
diff_hunk: @@ -237,6 +237,63 @@ def convert_warp_coords_to_mj_kernel(
             qvel[worldid, qd_i + i] = joint_qd[wqd_i + i]
 
 
+@wp.kernel
+def convert_joint_q_qpos0(
path: newton/solvers/solver_mujoco.py
user: login: adenzler-nvidia
body: I'm leaving this kernel in here although it's not called yet - something with qpos0 init is not working at all and breaks the simulation, but this needs thorough debugging and testing. I verified that this kernel is correct so it seems to be an issue with how we setup the sim, so we can still use this once we figure out how to set qpos0 correctly.
created_at: 2025-05-22T08:43:16Z
author_association: MEMBER
diff_hunk: @@ -1299,41 +1451,183 @@ def add_geoms(warp_body_id: int, perm_position=False, incoming_xform=None):
                 f.write(spec.to_xml())
                 print(f"Saved mujoco model to {os.path.abspath(target_filename)}")
 
-        d = mujoco.MjData(m)
+        self.mj_data = mujoco.MjData(self.mj_model)
 
-        d.nefc = nefc_per_env
+        self.mj_data.nefc = nefc_per_env
 
-        m.opt.tolerance = tolerance
-        m.opt.ls_tolerance = ls_tolerance
-        m.opt.cone = cone
-        m.opt.iterations = iterations
-        m.opt.ls_iterations = ls_iterations
-        m.opt.integrator = integrator
-        m.opt.solver = solver
+        self.mj_model.opt.tolerance = tolerance
+        self.mj_model.opt.ls_tolerance = ls_tolerance
+        self.mj_model.opt.cone = cone
+        self.mj_model.opt.iterations = iterations
+        self.mj_model.opt.ls_iterations = ls_iterations
+        self.mj_model.opt.integrator = integrator
+        self.mj_model.opt.solver = solver
         # m.opt.disableflags = disableflags
-        m.opt.impratio = impratio
-        m.opt.jacobian = mujoco.mjtJacobian.mjJAC_AUTO
+        self.mj_model.opt.impratio = impratio
+        self.mj_model.opt.jacobian = mujoco.mjtJacobian.mjJAC_AUTO
+
+        MuJoCoSolver.update_mjc_data(self.mj_data, model, state)
 
-        MuJoCoSolver.update_mjc_data(d, model, state)
+        # fill some MjWarp model fields that outdated after update_mjc_data.
+        # just setting qpos0 to d.qpos leads to weird behavior here, needs
+        # to be investigated.
 
-        mujoco.mj_forward(m, d)
+        mujoco.mj_forward(self.mj_model, self.mj_data)
 
         with wp.ScopedDevice(model.device):
             # add axis_to_actuator mapping to the Newton model
             model.axis_to_actuator = wp.array(axis_to_actuator, dtype=wp.int32)  # pyright: ignore[reportAttributeAccessIssue]
 
-            mj_model = mujoco_warp.put_model(m)
+            self.mjw_model = mujoco_warp.put_model(self.mj_model)
             if separate_envs_to_worlds:
                 nworld = model.num_envs
             else:
                 nworld = 1
+
+            # expand model fields that can be expanded:
+            self.expand_model_fields(self.mjw_model, nworld)
+
+            # complete the body mapping
+            size = max(body_mapping.keys()) + 1
+            arr = np.full(size, -1, dtype=int)
+            for k, v in body_mapping.items():
+                if k != -1:
+                    arr[k] = v
+
+            self.body_mapping = wp.array(arr, dtype=int)
+
+            # now fill with all the data from the Newton model.
+            flags = types.NOTIFY_FLAG_BODY_INERTIAL_PROPERTIES
+            self.notify_model_changed(flags)
+
             # TODO find better heuristics to determine nconmax and njmax
             if ncon_per_env:
                 nconmax = nworld * ncon_per_env
             else:
                 nconmax = model.rigid_contact_max * 4
-            nconmax = max(nconmax, d.ncon)
-            njmax = max(nworld * nefc_per_env * 4, nworld * d.nefc)
-            mj_data = mujoco_warp.put_data(m, d, nworld=nworld, nconmax=nconmax, njmax=njmax)
+            nconmax = max(nconmax, self.mj_data.ncon)
+            njmax = max(nworld * nefc_per_env * 4, nworld * self.mj_data.nefc)
+            self.mjw_data = mujoco_warp.put_data(
+                self.mj_model, self.mj_data, nworld=nworld, nconmax=nconmax, njmax=njmax
+            )
+
+    def expand_model_fields(self, mj_model: MjWarpModel, nworld: int):
+        if nworld == 1:
+            return
 
-        return mj_model, mj_data, m, d
+        model_fields_to_expand = [
path: newton/solvers/solver_mujoco.py
user: login: eric-heiden
body: So mujoco_warp just instantiates these fields for only one env and we want to explicitly replicate them here so that Newton can change body masses etc. for all envs independently, correct? Are these all possibly fields that can be expanded or just some subset that is supported for explicit replication in MJWarp? Maybe we can trim this list down for the few attributes we support for now?
created_at: 2025-05-22T17:44:01Z
author_association: MEMBER
diff_hunk: @@ -1299,41 +1451,183 @@ def add_geoms(warp_body_id: int, perm_position=False, incoming_xform=None):
                 f.write(spec.to_xml())
                 print(f"Saved mujoco model to {os.path.abspath(target_filename)}")
 
-        d = mujoco.MjData(m)
+        self.mj_data = mujoco.MjData(self.mj_model)
 
-        d.nefc = nefc_per_env
+        self.mj_data.nefc = nefc_per_env
 
-        m.opt.tolerance = tolerance
-        m.opt.ls_tolerance = ls_tolerance
-        m.opt.cone = cone
-        m.opt.iterations = iterations
-        m.opt.ls_iterations = ls_iterations
-        m.opt.integrator = integrator
-        m.opt.solver = solver
+        self.mj_model.opt.tolerance = tolerance
+        self.mj_model.opt.ls_tolerance = ls_tolerance
+        self.mj_model.opt.cone = cone
+        self.mj_model.opt.iterations = iterations
+        self.mj_model.opt.ls_iterations = ls_iterations
+        self.mj_model.opt.integrator = integrator
+        self.mj_model.opt.solver = solver
         # m.opt.disableflags = disableflags
-        m.opt.impratio = impratio
-        m.opt.jacobian = mujoco.mjtJacobian.mjJAC_AUTO
+        self.mj_model.opt.impratio = impratio
+        self.mj_model.opt.jacobian = mujoco.mjtJacobian.mjJAC_AUTO
+
+        MuJoCoSolver.update_mjc_data(self.mj_data, model, state)
 
-        MuJoCoSolver.update_mjc_data(d, model, state)
+        # fill some MjWarp model fields that outdated after update_mjc_data.
+        # just setting qpos0 to d.qpos leads to weird behavior here, needs
+        # to be investigated.
 
-        mujoco.mj_forward(m, d)
+        mujoco.mj_forward(self.mj_model, self.mj_data)
 
         with wp.ScopedDevice(model.device):
             # add axis_to_actuator mapping to the Newton model
             model.axis_to_actuator = wp.array(axis_to_actuator, dtype=wp.int32)  # pyright: ignore[reportAttributeAccessIssue]
 
-            mj_model = mujoco_warp.put_model(m)
+            self.mjw_model = mujoco_warp.put_model(self.mj_model)
             if separate_envs_to_worlds:
                 nworld = model.num_envs
             else:
                 nworld = 1
+
+            # expand model fields that can be expanded:
+            self.expand_model_fields(self.mjw_model, nworld)
+
+            # complete the body mapping
+            size = max(body_mapping.keys()) + 1
+            arr = np.full(size, -1, dtype=int)
+            for k, v in body_mapping.items():
+                if k != -1:
+                    arr[k] = v
+
+            self.body_mapping = wp.array(arr, dtype=int)
+
+            # now fill with all the data from the Newton model.
+            flags = types.NOTIFY_FLAG_BODY_INERTIAL_PROPERTIES
+            self.notify_model_changed(flags)
+
             # TODO find better heuristics to determine nconmax and njmax
             if ncon_per_env:
                 nconmax = nworld * ncon_per_env
             else:
                 nconmax = model.rigid_contact_max * 4
-            nconmax = max(nconmax, d.ncon)
-            njmax = max(nworld * nefc_per_env * 4, nworld * d.nefc)
-            mj_data = mujoco_warp.put_data(m, d, nworld=nworld, nconmax=nconmax, njmax=njmax)
+            nconmax = max(nconmax, self.mj_data.ncon)
+            njmax = max(nworld * nefc_per_env * 4, nworld * self.mj_data.nefc)
+            self.mjw_data = mujoco_warp.put_data(
+                self.mj_model, self.mj_data, nworld=nworld, nconmax=nconmax, njmax=njmax
+            )
+
+    def expand_model_fields(self, mj_model: MjWarpModel, nworld: int):
+        if nworld == 1:
+            return
 
-        return mj_model, mj_data, m, d
+        model_fields_to_expand = [
+            "qpos0",
+            "qpos_spring",
+            "body_pos",
+            "body_quat",
+            "body_ipos",
+            "body_iquat",
+            "body_mass",
+            "body_subtreemass",
+            "subtree_mass",
path: newton/solvers/solver_mujoco.py
user: login: eric-heiden
body: Do we need to update `subtree_mass` when updating body masses?
created_at: 2025-05-22T17:48:26Z
author_association: MEMBER
diff_hunk: @@ -1299,41 +1451,183 @@ def add_geoms(warp_body_id: int, perm_position=False, incoming_xform=None):
                 f.write(spec.to_xml())
                 print(f"Saved mujoco model to {os.path.abspath(target_filename)}")
 
-        d = mujoco.MjData(m)
+        self.mj_data = mujoco.MjData(self.mj_model)
 
-        d.nefc = nefc_per_env
+        self.mj_data.nefc = nefc_per_env
 
-        m.opt.tolerance = tolerance
-        m.opt.ls_tolerance = ls_tolerance
-        m.opt.cone = cone
-        m.opt.iterations = iterations
-        m.opt.ls_iterations = ls_iterations
-        m.opt.integrator = integrator
-        m.opt.solver = solver
+        self.mj_model.opt.tolerance = tolerance
+        self.mj_model.opt.ls_tolerance = ls_tolerance
+        self.mj_model.opt.cone = cone
+        self.mj_model.opt.iterations = iterations
+        self.mj_model.opt.ls_iterations = ls_iterations
+        self.mj_model.opt.integrator = integrator
+        self.mj_model.opt.solver = solver
         # m.opt.disableflags = disableflags
-        m.opt.impratio = impratio
-        m.opt.jacobian = mujoco.mjtJacobian.mjJAC_AUTO
+        self.mj_model.opt.impratio = impratio
+        self.mj_model.opt.jacobian = mujoco.mjtJacobian.mjJAC_AUTO
+
+        MuJoCoSolver.update_mjc_data(self.mj_data, model, state)
 
-        MuJoCoSolver.update_mjc_data(d, model, state)
+        # fill some MjWarp model fields that outdated after update_mjc_data.
+        # just setting qpos0 to d.qpos leads to weird behavior here, needs
+        # to be investigated.
 
-        mujoco.mj_forward(m, d)
+        mujoco.mj_forward(self.mj_model, self.mj_data)
 
         with wp.ScopedDevice(model.device):
             # add axis_to_actuator mapping to the Newton model
             model.axis_to_actuator = wp.array(axis_to_actuator, dtype=wp.int32)  # pyright: ignore[reportAttributeAccessIssue]
 
-            mj_model = mujoco_warp.put_model(m)
+            self.mjw_model = mujoco_warp.put_model(self.mj_model)
             if separate_envs_to_worlds:
                 nworld = model.num_envs
             else:
                 nworld = 1
+
+            # expand model fields that can be expanded:
+            self.expand_model_fields(self.mjw_model, nworld)
+
+            # complete the body mapping
+            size = max(body_mapping.keys()) + 1
+            arr = np.full(size, -1, dtype=int)
+            for k, v in body_mapping.items():
+                if k != -1:
+                    arr[k] = v
+
+            self.body_mapping = wp.array(arr, dtype=int)
+
+            # now fill with all the data from the Newton model.
+            flags = types.NOTIFY_FLAG_BODY_INERTIAL_PROPERTIES
+            self.notify_model_changed(flags)
+
             # TODO find better heuristics to determine nconmax and njmax
             if ncon_per_env:
                 nconmax = nworld * ncon_per_env
             else:
                 nconmax = model.rigid_contact_max * 4
-            nconmax = max(nconmax, d.ncon)
-            njmax = max(nworld * nefc_per_env * 4, nworld * d.nefc)
-            mj_data = mujoco_warp.put_data(m, d, nworld=nworld, nconmax=nconmax, njmax=njmax)
+            nconmax = max(nconmax, self.mj_data.ncon)
+            njmax = max(nworld * nefc_per_env * 4, nworld * self.mj_data.nefc)
+            self.mjw_data = mujoco_warp.put_data(
+                self.mj_model, self.mj_data, nworld=nworld, nconmax=nconmax, njmax=njmax
+            )
+
+    def expand_model_fields(self, mj_model: MjWarpModel, nworld: int):
+        if nworld == 1:
+            return
 
-        return mj_model, mj_data, m, d
+        model_fields_to_expand = [
+            "qpos0",
+            "qpos_spring",
+            "body_pos",
+            "body_quat",
+            "body_ipos",
+            "body_iquat",
+            "body_mass",
+            "body_subtreemass",
+            "subtree_mass",
path: newton/solvers/solver_mujoco.py
user: login: eric-heiden
body: I'm also not sure if MJWarp can handle the different dimensions - it looks like `smooth.com_pos` assumes `model.subtree_mass` and `data.subtree_com` have the same shape, right? Or is there some zero-stride trick being used?
created_at: 2025-05-22T17:49:57Z
author_association: MEMBER
diff_hunk: @@ -1163,6 +1314,7 @@ def add_geoms(warp_body_id: int, perm_position=False, incoming_xform=None):
                 mass=body_mass[child],
                 ipos=pos2mjc(body_com[child, :]),
                 fullinertia=[inertia[0, 0], inertia[1, 1], inertia[2, 2], inertia[0, 1], inertia[0, 2], inertia[1, 2]],
+                explicitinertial=True,
path: newton/solvers/solver_mujoco.py
user: login: eric-heiden
body: This is great! Thanks for fixing this, this should write the inertia to the MJCF hopefully. Although we might have to double-check the frame / world up axis handling here potentially.

Let's link this to https://github.com/newton-physics/newton/issues/125
created_at: 2025-05-22T17:54:07Z
author_association: MEMBER
diff_hunk: @@ -347,6 +347,13 @@ def __hash__(self) -> int:
         return hash((tuple(np.array(self.vertices).flatten()), tuple(np.array(self.indices).flatten()), self.is_solid))
 
 
+# model update flags - WIP grouping and naming still being worked on.
+NOTIFY_FLAG_JOINT_PROPERTIES = wp.constant(1 << 0)
+NOTIFY_FLAG_DOF_PROPERTIES = wp.constant(1 << 1)
+NOTIFY_FLAG_BODY_PROPERTIES = wp.constant(1 << 2)
+NOTIFY_FLAG_BODY_INERTIAL_PROPERTIES = wp.constant(1 << 3)
+NOTIFY_FLAG_SHAPE_PROPERTIES = wp.constant(1 << 4)
path: newton/core/types.py
user: login: eric-heiden
body: I think this is a good start! It would be nice to clarify via some docstrings what each of these flags exactly mean, I'm not sure what the differences are between `NOTIFY_FLAG_BODY_PROPERTIES`  and `NOTIFY_FLAG_BODY_INERTIAL_PROPERTIES`, and what is meant by `NOTIFY_FLAG_DOF_PROPERTIES`. Ideally there shouldn't be any overlap between these flags.
created_at: 2025-05-22T17:58:28Z
author_association: MEMBER
diff_hunk: @@ -347,6 +347,13 @@ def __hash__(self) -> int:
         return hash((tuple(np.array(self.vertices).flatten()), tuple(np.array(self.indices).flatten()), self.is_solid))
 
 
+# model update flags - WIP grouping and naming still being worked on.
+NOTIFY_FLAG_JOINT_PROPERTIES = wp.constant(1 << 0)
+NOTIFY_FLAG_DOF_PROPERTIES = wp.constant(1 << 1)
+NOTIFY_FLAG_BODY_PROPERTIES = wp.constant(1 << 2)
+NOTIFY_FLAG_BODY_INERTIAL_PROPERTIES = wp.constant(1 << 3)
+NOTIFY_FLAG_SHAPE_PROPERTIES = wp.constant(1 << 4)
path: newton/core/types.py
user: login: adenzler-nvidia
body: I can add some comments. I would say let's iterate on these while we go, and then clean up once we are closer to MVP/have support for every model field we can and want to be randomizable.
created_at: 2025-05-22T18:26:17Z
author_association: MEMBER
diff_hunk: @@ -1299,41 +1451,183 @@ def add_geoms(warp_body_id: int, perm_position=False, incoming_xform=None):
                 f.write(spec.to_xml())
                 print(f"Saved mujoco model to {os.path.abspath(target_filename)}")
 
-        d = mujoco.MjData(m)
+        self.mj_data = mujoco.MjData(self.mj_model)
 
-        d.nefc = nefc_per_env
+        self.mj_data.nefc = nefc_per_env
 
-        m.opt.tolerance = tolerance
-        m.opt.ls_tolerance = ls_tolerance
-        m.opt.cone = cone
-        m.opt.iterations = iterations
-        m.opt.ls_iterations = ls_iterations
-        m.opt.integrator = integrator
-        m.opt.solver = solver
+        self.mj_model.opt.tolerance = tolerance
+        self.mj_model.opt.ls_tolerance = ls_tolerance
+        self.mj_model.opt.cone = cone
+        self.mj_model.opt.iterations = iterations
+        self.mj_model.opt.ls_iterations = ls_iterations
+        self.mj_model.opt.integrator = integrator
+        self.mj_model.opt.solver = solver
         # m.opt.disableflags = disableflags
-        m.opt.impratio = impratio
-        m.opt.jacobian = mujoco.mjtJacobian.mjJAC_AUTO
+        self.mj_model.opt.impratio = impratio
+        self.mj_model.opt.jacobian = mujoco.mjtJacobian.mjJAC_AUTO
+
+        MuJoCoSolver.update_mjc_data(self.mj_data, model, state)
 
-        MuJoCoSolver.update_mjc_data(d, model, state)
+        # fill some MjWarp model fields that outdated after update_mjc_data.
+        # just setting qpos0 to d.qpos leads to weird behavior here, needs
+        # to be investigated.
 
-        mujoco.mj_forward(m, d)
+        mujoco.mj_forward(self.mj_model, self.mj_data)
 
         with wp.ScopedDevice(model.device):
             # add axis_to_actuator mapping to the Newton model
             model.axis_to_actuator = wp.array(axis_to_actuator, dtype=wp.int32)  # pyright: ignore[reportAttributeAccessIssue]
 
-            mj_model = mujoco_warp.put_model(m)
+            self.mjw_model = mujoco_warp.put_model(self.mj_model)
             if separate_envs_to_worlds:
                 nworld = model.num_envs
             else:
                 nworld = 1
+
+            # expand model fields that can be expanded:
+            self.expand_model_fields(self.mjw_model, nworld)
+
+            # complete the body mapping
+            size = max(body_mapping.keys()) + 1
+            arr = np.full(size, -1, dtype=int)
+            for k, v in body_mapping.items():
+                if k != -1:
+                    arr[k] = v
+
+            self.body_mapping = wp.array(arr, dtype=int)
+
+            # now fill with all the data from the Newton model.
+            flags = types.NOTIFY_FLAG_BODY_INERTIAL_PROPERTIES
+            self.notify_model_changed(flags)
+
             # TODO find better heuristics to determine nconmax and njmax
             if ncon_per_env:
                 nconmax = nworld * ncon_per_env
             else:
                 nconmax = model.rigid_contact_max * 4
-            nconmax = max(nconmax, d.ncon)
-            njmax = max(nworld * nefc_per_env * 4, nworld * d.nefc)
-            mj_data = mujoco_warp.put_data(m, d, nworld=nworld, nconmax=nconmax, njmax=njmax)
+            nconmax = max(nconmax, self.mj_data.ncon)
+            njmax = max(nworld * nefc_per_env * 4, nworld * self.mj_data.nefc)
+            self.mjw_data = mujoco_warp.put_data(
+                self.mj_model, self.mj_data, nworld=nworld, nconmax=nconmax, njmax=njmax
+            )
+
+    def expand_model_fields(self, mj_model: MjWarpModel, nworld: int):
+        if nworld == 1:
+            return
 
-        return mj_model, mj_data, m, d
+        model_fields_to_expand = [
path: newton/solvers/solver_mujoco.py
user: login: adenzler-nvidia
body: mujoco warp uses an additional dimension with stride zero for all of these - which can then be replaced by expanded arrays as you see fit. 

The list contains everything that is using stride 0, but we can trim it down to only the ones that have a corresponding newton model field. Maybe we might want to expose some of them as mujoco specific fields in the future but that can wait for sure.

I pasted the whole list to keep an overview, maybe we just comment out all the ones that are not implemented, and then remove the commented ones that we don't care about after that?
created_at: 2025-05-22T18:29:01Z
author_association: MEMBER
diff_hunk: @@ -1299,41 +1451,183 @@ def add_geoms(warp_body_id: int, perm_position=False, incoming_xform=None):
                 f.write(spec.to_xml())
                 print(f"Saved mujoco model to {os.path.abspath(target_filename)}")
 
-        d = mujoco.MjData(m)
+        self.mj_data = mujoco.MjData(self.mj_model)
 
-        d.nefc = nefc_per_env
+        self.mj_data.nefc = nefc_per_env
 
-        m.opt.tolerance = tolerance
-        m.opt.ls_tolerance = ls_tolerance
-        m.opt.cone = cone
-        m.opt.iterations = iterations
-        m.opt.ls_iterations = ls_iterations
-        m.opt.integrator = integrator
-        m.opt.solver = solver
+        self.mj_model.opt.tolerance = tolerance
+        self.mj_model.opt.ls_tolerance = ls_tolerance
+        self.mj_model.opt.cone = cone
+        self.mj_model.opt.iterations = iterations
+        self.mj_model.opt.ls_iterations = ls_iterations
+        self.mj_model.opt.integrator = integrator
+        self.mj_model.opt.solver = solver
         # m.opt.disableflags = disableflags
-        m.opt.impratio = impratio
-        m.opt.jacobian = mujoco.mjtJacobian.mjJAC_AUTO
+        self.mj_model.opt.impratio = impratio
+        self.mj_model.opt.jacobian = mujoco.mjtJacobian.mjJAC_AUTO
+
+        MuJoCoSolver.update_mjc_data(self.mj_data, model, state)
 
-        MuJoCoSolver.update_mjc_data(d, model, state)
+        # fill some MjWarp model fields that outdated after update_mjc_data.
+        # just setting qpos0 to d.qpos leads to weird behavior here, needs
+        # to be investigated.
 
-        mujoco.mj_forward(m, d)
+        mujoco.mj_forward(self.mj_model, self.mj_data)
 
         with wp.ScopedDevice(model.device):
             # add axis_to_actuator mapping to the Newton model
             model.axis_to_actuator = wp.array(axis_to_actuator, dtype=wp.int32)  # pyright: ignore[reportAttributeAccessIssue]
 
-            mj_model = mujoco_warp.put_model(m)
+            self.mjw_model = mujoco_warp.put_model(self.mj_model)
             if separate_envs_to_worlds:
                 nworld = model.num_envs
             else:
                 nworld = 1
+
+            # expand model fields that can be expanded:
+            self.expand_model_fields(self.mjw_model, nworld)
+
+            # complete the body mapping
+            size = max(body_mapping.keys()) + 1
+            arr = np.full(size, -1, dtype=int)
+            for k, v in body_mapping.items():
+                if k != -1:
+                    arr[k] = v
+
+            self.body_mapping = wp.array(arr, dtype=int)
+
+            # now fill with all the data from the Newton model.
+            flags = types.NOTIFY_FLAG_BODY_INERTIAL_PROPERTIES
+            self.notify_model_changed(flags)
+
             # TODO find better heuristics to determine nconmax and njmax
             if ncon_per_env:
                 nconmax = nworld * ncon_per_env
             else:
                 nconmax = model.rigid_contact_max * 4
-            nconmax = max(nconmax, d.ncon)
-            njmax = max(nworld * nefc_per_env * 4, nworld * d.nefc)
-            mj_data = mujoco_warp.put_data(m, d, nworld=nworld, nconmax=nconmax, njmax=njmax)
+            nconmax = max(nconmax, self.mj_data.ncon)
+            njmax = max(nworld * nefc_per_env * 4, nworld * self.mj_data.nefc)
+            self.mjw_data = mujoco_warp.put_data(
+                self.mj_model, self.mj_data, nworld=nworld, nconmax=nconmax, njmax=njmax
+            )
+
+    def expand_model_fields(self, mj_model: MjWarpModel, nworld: int):
+        if nworld == 1:
+            return
 
-        return mj_model, mj_data, m, d
+        model_fields_to_expand = [
+            "qpos0",
+            "qpos_spring",
+            "body_pos",
+            "body_quat",
+            "body_ipos",
+            "body_iquat",
+            "body_mass",
+            "body_subtreemass",
+            "subtree_mass",
path: newton/solvers/solver_mujoco.py
user: login: adenzler-nvidia
body: These are using stride 0 in the first dimension, as stated above. So we're all good on the implementation side.

There are some dependencies between the fields, we will have to discover them as we go. I also still don't understand the difference between subtree_mass and body_subtreemass. I think it's futile to do this upfront but we will discover as we implement.
created_at: 2025-05-22T18:30:29Z
author_association: MEMBER
diff_hunk: @@ -1299,41 +1451,183 @@ def add_geoms(warp_body_id: int, perm_position=False, incoming_xform=None):
                 f.write(spec.to_xml())
                 print(f"Saved mujoco model to {os.path.abspath(target_filename)}")
 
-        d = mujoco.MjData(m)
+        self.mj_data = mujoco.MjData(self.mj_model)
 
-        d.nefc = nefc_per_env
+        self.mj_data.nefc = nefc_per_env
 
-        m.opt.tolerance = tolerance
-        m.opt.ls_tolerance = ls_tolerance
-        m.opt.cone = cone
-        m.opt.iterations = iterations
-        m.opt.ls_iterations = ls_iterations
-        m.opt.integrator = integrator
-        m.opt.solver = solver
+        self.mj_model.opt.tolerance = tolerance
+        self.mj_model.opt.ls_tolerance = ls_tolerance
+        self.mj_model.opt.cone = cone
+        self.mj_model.opt.iterations = iterations
+        self.mj_model.opt.ls_iterations = ls_iterations
+        self.mj_model.opt.integrator = integrator
+        self.mj_model.opt.solver = solver
         # m.opt.disableflags = disableflags
-        m.opt.impratio = impratio
-        m.opt.jacobian = mujoco.mjtJacobian.mjJAC_AUTO
+        self.mj_model.opt.impratio = impratio
+        self.mj_model.opt.jacobian = mujoco.mjtJacobian.mjJAC_AUTO
+
+        MuJoCoSolver.update_mjc_data(self.mj_data, model, state)
 
-        MuJoCoSolver.update_mjc_data(d, model, state)
+        # fill some MjWarp model fields that outdated after update_mjc_data.
+        # just setting qpos0 to d.qpos leads to weird behavior here, needs
+        # to be investigated.
 
-        mujoco.mj_forward(m, d)
+        mujoco.mj_forward(self.mj_model, self.mj_data)
 
         with wp.ScopedDevice(model.device):
             # add axis_to_actuator mapping to the Newton model
             model.axis_to_actuator = wp.array(axis_to_actuator, dtype=wp.int32)  # pyright: ignore[reportAttributeAccessIssue]
 
-            mj_model = mujoco_warp.put_model(m)
+            self.mjw_model = mujoco_warp.put_model(self.mj_model)
             if separate_envs_to_worlds:
                 nworld = model.num_envs
             else:
                 nworld = 1
+
+            # expand model fields that can be expanded:
+            self.expand_model_fields(self.mjw_model, nworld)
+
+            # complete the body mapping
+            size = max(body_mapping.keys()) + 1
+            arr = np.full(size, -1, dtype=int)
+            for k, v in body_mapping.items():
+                if k != -1:
+                    arr[k] = v
+
+            self.body_mapping = wp.array(arr, dtype=int)
+
+            # now fill with all the data from the Newton model.
+            flags = types.NOTIFY_FLAG_BODY_INERTIAL_PROPERTIES
+            self.notify_model_changed(flags)
+
             # TODO find better heuristics to determine nconmax and njmax
             if ncon_per_env:
                 nconmax = nworld * ncon_per_env
             else:
                 nconmax = model.rigid_contact_max * 4
-            nconmax = max(nconmax, d.ncon)
-            njmax = max(nworld * nefc_per_env * 4, nworld * d.nefc)
-            mj_data = mujoco_warp.put_data(m, d, nworld=nworld, nconmax=nconmax, njmax=njmax)
+            nconmax = max(nconmax, self.mj_data.ncon)
+            njmax = max(nworld * nefc_per_env * 4, nworld * self.mj_data.nefc)
+            self.mjw_data = mujoco_warp.put_data(
+                self.mj_model, self.mj_data, nworld=nworld, nconmax=nconmax, njmax=njmax
+            )
+
+    def expand_model_fields(self, mj_model: MjWarpModel, nworld: int):
+        if nworld == 1:
+            return
 
-        return mj_model, mj_data, m, d
+        model_fields_to_expand = [
path: newton/solvers/solver_mujoco.py
user: login: eric-heiden
body: Sounds good, maybe yeah commenting out the fields we are not yet updating would be useful.
created_at: 2025-05-22T19:38:33Z
author_association: MEMBER
diff_hunk: @@ -237,6 +237,63 @@ def convert_warp_coords_to_mj_kernel(
             qvel[worldid, qd_i + i] = joint_qd[wqd_i + i]
 
 
+@wp.kernel
+def convert_joint_q_qpos0(
path: newton/solvers/solver_mujoco.py
user: login: eric-heiden
body: It would be good to have some more code reuse here since there are kernels already converting qpos, qvel to Newton. So if it's not working anyway we can just remove `convert_joint_q_qpos0`?
created_at: 2025-05-22T19:39:38Z
author_association: MEMBER
diff_hunk: @@ -347,6 +347,13 @@ def __hash__(self) -> int:
         return hash((tuple(np.array(self.vertices).flatten()), tuple(np.array(self.indices).flatten()), self.is_solid))
 
 
+# model update flags - WIP grouping and naming still being worked on.
+NOTIFY_FLAG_JOINT_PROPERTIES = wp.constant(1 << 0)
+NOTIFY_FLAG_DOF_PROPERTIES = wp.constant(1 << 1)
+NOTIFY_FLAG_BODY_PROPERTIES = wp.constant(1 << 2)
+NOTIFY_FLAG_BODY_INERTIAL_PROPERTIES = wp.constant(1 << 3)
+NOTIFY_FLAG_SHAPE_PROPERTIES = wp.constant(1 << 4)
path: newton/core/types.py
user: login: eric-heiden
body: For sure, we can iterate on this in later PR's. Some comments would be nice to get the ball rolling.
created_at: 2025-05-22T19:40:20Z
author_association: MEMBER
diff_hunk: @@ -347,6 +347,13 @@ def __hash__(self) -> int:
         return hash((tuple(np.array(self.vertices).flatten()), tuple(np.array(self.indices).flatten()), self.is_solid))
 
 
+# model update flags - WIP grouping and naming still being worked on.
+NOTIFY_FLAG_JOINT_PROPERTIES = wp.constant(1 << 0)
+NOTIFY_FLAG_DOF_PROPERTIES = wp.constant(1 << 1)
+NOTIFY_FLAG_BODY_PROPERTIES = wp.constant(1 << 2)
+NOTIFY_FLAG_BODY_INERTIAL_PROPERTIES = wp.constant(1 << 3)
+NOTIFY_FLAG_SHAPE_PROPERTIES = wp.constant(1 << 4)
path: newton/core/types.py
user: login: adenzler-nvidia
body: extended the list and added the model properties that I think make sense in the comments.
created_at: 2025-05-23T08:45:42Z
author_association: MEMBER
diff_hunk: @@ -237,6 +237,63 @@ def convert_warp_coords_to_mj_kernel(
             qvel[worldid, qd_i + i] = joint_qd[wqd_i + i]
 
 
+@wp.kernel
+def convert_joint_q_qpos0(
path: newton/solvers/solver_mujoco.py
user: login: adenzler-nvidia
body: I can remove it yes
created_at: 2025-05-23T08:46:28Z
author_association: MEMBER

PR comments:
user: login: adenzler-nvidia
created_at: 2025-05-19T19:33:40Z
body: moving to draft as there are too many moving pieces I need to connect first here.
author_association: MEMBER
user: login: eric-heiden
created_at: 2025-05-22T19:37:58Z
body: In my opinion this is a good first step to get the change notification system going and have a stab at updating the MuJoCo model from Newton. The inertial property setting is also better to have now. So I would say we should merge this soon and iterate from here.
author_association: MEMBER