Fix State.assign for namespaced extended and custom attributes

CHANGELOG.md

@@ -35,6 +35,7 @@
 - Show prismatic joints in the GL viewer when "Show Joints" is enabled
 - Fix connect constraint anchor computation to account for joint reference positions when `SolverMuJoCo` is the chosen solver.
 - Fix `SolverMuJoCo` passing non-zero geom/pair margins to `mujoco_warp.put_model()`, which fails when NATIVECCD is enabled. Margins are forced to zero when MuJoCo handles collisions (`use_mujoco_contacts=True`); the Newton collision pipeline (`use_mujoco_contacts=False`) is unchanged
+- Fix `State.assign` not copying namespaced extended and custom state attributes 
 - Fix mesh-convex back-face contacts generating inverted normals that trap shapes inside meshes and cause solver divergence (NaN)
 - Fix finite plane geometry 2x too large in collision, bounding sphere, and raytrace sensor
 - Fix MPR convergence failure on large and extreme-aspect-ratio mesh triangles by projecting the starting point onto the triangle nearest the convex center

newton/_src/sim/state.py

@@ -6,6 +6,44 @@
 import warp as wp
 
 
+def _copy_arrays(dst: object, src: object, prefix: str = "") -> None:
+    """Copy ``wp.array`` attributes from ``src`` into ``dst``.
+
+    Walks both objects' ``__dict__``, matches attributes by name, and copies
+    ``wp.array`` values via ``dst_array.assign(src_array)``. Raises
+    :class:`ValueError` on presence mismatch (one side has an array where
+    the other does not).
+
+    Args:
+        dst: Destination object (its ``wp.array`` attributes will be overwritten).
+        src: Source object.
+        prefix: Prefix prepended to attribute names in error messages
+            (e.g. ``"mujoco."`` for namespaced attributes).
+    """
+    attributes = set(dst.__dict__).union(src.__dict__)
+    for attr in attributes:
+        val_dst = getattr(dst, attr, None)
+        val_src = getattr(src, attr, None)
+
+        if val_dst is None and val_src is None:
+            continue
+
+        array_dst = isinstance(val_dst, wp.array)
+        array_src = isinstance(val_src, wp.array)
+
+        if not array_dst and not array_src:
+            continue
+
+        qualified = f"{prefix}{attr}"
+        if val_dst is None or not array_dst:
+            raise ValueError(f"State is missing array for '{qualified}' which is present in the other state.")
+
+        if val_src is None or not array_src:
+            raise ValueError(f"Other state is missing array for '{qualified}' which is present in this state.")
+
+        val_dst.assign(val_src)
+
+
 class State:
     """
     Represents the time-varying state of a :class:`Model` in a simulation.
@@ -157,28 +195,38 @@ def assign(self, other: State) -> None:
         Raises:
             ValueError: If the states have mismatched attributes (one has an array allocated where the other is None).
         """
-        attributes = set(self.__dict__).union(other.__dict__)
-
-        for attr in attributes:
-            val_self = getattr(self, attr, None)
-            val_other = getattr(other, attr, None)
-
-            if val_self is None and val_other is None:
+        from .model import Model  # noqa: PLC0415
+
+        # Top-level array attributes.
+        _copy_arrays(self, other)
+
+        # Discover all AttributeNamespace containers on either state and
+        # descend into each. This uniformly covers both EXTENDED_ATTRIBUTES
+        # (e.g. ``mujoco.qfrc_actuator``) and custom namespaced attributes
+        # registered via ``ModelBuilder.add_custom_attribute``.
+        ns_self = {k: v for k, v in self.__dict__.items() if isinstance(v, Model.AttributeNamespace)}
+        ns_other = {k: v for k, v in other.__dict__.items() if isinstance(v, Model.AttributeNamespace)}
+
+        for ns_name in ns_self.keys() | ns_other.keys():
+            dst_ns = ns_self.get(ns_name)
+            src_ns = ns_other.get(ns_name)
+
+            # If the namespace container is missing on one side, only raise
+            # when the other side actually holds arrays inside it.
+            if dst_ns is None:
+                if any(isinstance(v, wp.array) for v in src_ns.__dict__.values()):
+                    raise ValueError(
+                        f"State is missing namespace '{ns_name}' which contains arrays in the other state."
+                    )
                 continue
-
-            array_self = isinstance(val_self, wp.array)
-            array_other = isinstance(val_other, wp.array)
-
-            if not array_self and not array_other:
+            if src_ns is None:
+                if any(isinstance(v, wp.array) for v in dst_ns.__dict__.values()):
+                    raise ValueError(
+                        f"Other state is missing namespace '{ns_name}' which contains arrays in this state."
+                    )
                 continue
 
-            if val_self is None or not array_self:
-                raise ValueError(f"State is missing array for '{attr}' which is present in the other state.")
-
-            if val_other is None or not array_other:
-                raise ValueError(f"Other state is missing array for '{attr}' which is present in this state.")
-
-            val_self.assign(val_other)
+            _copy_arrays(dst_ns, src_ns, prefix=f"{ns_name}.")
 
     @property
     def requires_grad(self) -> bool:

newton/tests/test_state_assign.py

@@ -0,0 +1,135 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 The Newton Developers
+# SPDX-License-Identifier: Apache-2.0
+
+import unittest
+
+import numpy as np
+import warp as wp
+
+import newton
+
+
+def _build_model(*, custom_attrs: tuple[str, ...] = ()):
+    builder = newton.ModelBuilder(up_axis=newton.Axis.Y, gravity=0.0)
+    inertia = wp.mat33((0.1, 0.0, 0.0), (0.0, 0.1, 0.0), (0.0, 0.0, 0.1))
+    body = builder.add_link(armature=0.0, inertia=inertia, mass=1.0)
+    joint = builder.add_joint_revolute(
+        parent=-1,
+        child=body,
+        parent_xform=wp.transform(wp.vec3(0.0, 0.0, 0.0), wp.quat_identity()),
+        child_xform=wp.transform(wp.vec3(0.0, 0.0, 0.0), wp.quat_identity()),
+        axis=wp.vec3(0.0, 0.0, 1.0),
+        target_pos=0.0,
+        target_ke=100.0,
+        target_kd=10.0,
+        effort_limit=5.0,
+        actuator_mode=newton.JointTargetMode.POSITION_VELOCITY,
+    )
+    builder.add_articulation([joint])
+    builder.request_state_attributes("mujoco:qfrc_actuator")
+    for name in custom_attrs:
+        builder.add_custom_attribute(
+            newton.ModelBuilder.CustomAttribute(
+                name=name,
+                frequency=newton.Model.AttributeFrequency.BODY,
+                dtype=wp.float32,
+                default=0.0,
+                assignment=newton.Model.AttributeAssignment.STATE,
+                namespace="my_namespace",
+            )
+        )
+    model = builder.finalize()
+    model.ground = False
+    return model
+
+
+class TestStateAssignNamespacedAttributes(unittest.TestCase):
+    def test_copies_namespaced_attribute(self):
+        model = _build_model()
+        state_0 = model.state()
+        state_1 = model.state()
+
+        sentinel = np.array([3.14], dtype=np.float32)
+        state_1.mujoco.qfrc_actuator.assign(sentinel)
+
+        state_0.assign(state_1)
+
+        np.testing.assert_allclose(state_0.mujoco.qfrc_actuator.numpy(), sentinel)
+
+    def test_raises_when_src_missing_namespaced_attribute(self):
+        model = _build_model()
+        state_0 = model.state()
+        state_1 = model.state()
+        delattr(state_1, "mujoco")
+
+        with self.assertRaises(ValueError):
+            state_0.assign(state_1)
+
+    def test_raises_when_dst_missing_namespaced_attribute(self):
+        model = _build_model()
+        state_0 = model.state()
+        state_1 = model.state()
+        delattr(state_0, "mujoco")
+
+        with self.assertRaises(ValueError):
+            state_0.assign(state_1)
+
+    def test_copies_custom_namespaced_attribute(self):
+        model = _build_model(custom_attrs=("my_attribute",))
+        state_0 = model.state()
+        state_1 = model.state()
+
+        sentinel = np.array([2.71], dtype=np.float32)
+        state_1.my_namespace.my_attribute.assign(sentinel)
+
+        state_0.assign(state_1)
+
+        np.testing.assert_allclose(state_0.my_namespace.my_attribute.numpy(), sentinel)
+
+    def test_raises_when_src_missing_custom_namespaced_attribute(self):
+        model = _build_model(custom_attrs=("my_attribute",))
+        state_0 = model.state()
+        state_1 = model.state()
+        delattr(state_1, "my_namespace")
+
+        with self.assertRaises(ValueError):
+            state_0.assign(state_1)
+
+    def test_raises_when_dst_missing_custom_namespaced_attribute(self):
+        model = _build_model(custom_attrs=("my_attribute",))
+        state_0 = model.state()
+        state_1 = model.state()
+        delattr(state_0, "my_namespace")
+
+        with self.assertRaises(ValueError):
+            state_0.assign(state_1)
+
+    def test_copies_multiple_custom_namespaced_attributes(self):
+        model = _build_model(custom_attrs=("attr_one", "attr_two"))
+        state_0 = model.state()
+        state_1 = model.state()
+
+        sentinel_one = np.array([1.23], dtype=np.float32)
+        sentinel_two = np.array([4.56], dtype=np.float32)
+        state_1.my_namespace.attr_one.assign(sentinel_one)
+        state_1.my_namespace.attr_two.assign(sentinel_two)
+
+        state_0.assign(state_1)
+
+        np.testing.assert_allclose(state_0.my_namespace.attr_one.numpy(), sentinel_one)
+        np.testing.assert_allclose(state_0.my_namespace.attr_two.numpy(), sentinel_two)
+
+    def test_raises_when_one_of_multiple_custom_attributes_missing(self):
+        model = _build_model(custom_attrs=("attr_one", "attr_two"))
+        state_0 = model.state()
+        state_1 = model.state()
+        # Remove a single attribute inside the namespace container (not the
+        # container itself) to exercise per-attribute presence checks.
+        delattr(state_1.my_namespace, "attr_two")
+
+        with self.assertRaises(ValueError):
+            state_0.assign(state_1)
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)