diff --git a/CHANGELOG.md b/CHANGELOG.md index f93523bbf0..24800bee81 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/newton/_src/sim/state.py b/newton/_src/sim/state.py index fa6efd9923..184c6c2a35 100644 --- a/newton/_src/sim/state.py +++ b/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: diff --git a/newton/tests/test_state_assign.py b/newton/tests/test_state_assign.py new file mode 100644 index 0000000000..6a89a64ac5 --- /dev/null +++ b/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)