Fix State.assign for namespaced extended and custom attributes

[jsw7460]

Description

State.assign skips wp.array attributes nested under attribute namespaces (e.g. state.mujoco.qfrc_actuator). The top-level loop only checks isinstance(val, wp.array) on __dict__ entries, so namespace containers are skipped and never descended into. After state_0.assign(state_1), state_0.mujoco.qfrc_actuator is still zero.

Fix: after the top-level loop, walk EXTENDED_ATTRIBUTES for "namespace:attr" entries and apply the same presence/mismatch checks.

Checklist

• New or existing tests cover these changes
• The documentation is up to date with these changes
• CHANGELOG.md has been updated (if user-facing change)

Test plan

uv run --extra dev -m newton.tests -k test_state_assign

Bug fix

Steps to reproduce:

1. Request mujoco:qfrc_actuator on a model and step SolverMuJoCo so state_1.mujoco.qfrc_actuator is populated.
2. Call state_0.assign(state_1) (the pattern State.assign's docstring recommends for odd-substep CUDA-graph loops).
3. Read state_0.mujoco.qfrc_actuator: values are still the zero-initialized array.

Minimal reproduction:

See test_state_assign.py

Note: I followed the current "namespace:attr" convention in EXTENDED_ATTRIBUTES.
If deeper nesting is needed in the future, let me know and I'll switch to a recursive approach.

Summary by CodeRabbit

• Bug Fixes

  • State assignment now correctly copies namespaced extended and custom state attributes between states and raises on mismatches.

• Tests

  • Added tests that confirm namespaced attributes are copied correctly and that assignment fails when required namespace data is missing.

• Refactor

  • Internal state-copy logic reorganized to more reliably handle top-level and namespaced array transfers.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Pro

Run ID: de89c8b9-59e8-4664-9648-0cbf5221a5a9

📥 Commits

Reviewing files that changed from the base of the PR and between 11ce9cf and 8fca42e.

📒 Files selected for processing (1)

• CHANGELOG.md

🚧 Files skipped from review as they are similar to previous changes (1)

• CHANGELOG.md

---

📝 Walkthrough

Walkthrough

Refactors State.assign() to copy wp.array attributes in two phases: top-level arrays and namespaced arrays via Model.AttributeNamespace containers, adding a helper to perform namespace-prefixed array copying and raising ValueError on per-namespace array-allocation mismatches.

Changes

Cohort / File(s)	Summary
State Transfer Logic newton/_src/sim/state.py	Reworked State.assign() to delegate array copying to a new private _copy_arrays helper. Copies top-level wp.array attributes and iterates Model.AttributeNamespace containers to copy namespace.attribute arrays with prefix-aware copying. Per-namespace validation now raises ValueError if arrays exist on one side but not the other; previous per-attribute mismatch checks for non-array fields removed.
Test Coverage newton/tests/test_state_assign.py	New unit tests validating namespaced attribute copying and error cases. Builds minimal model, requests built-in mujoco:qfrc_actuator and optional custom namespaced attributes, asserts correct array copying via .numpy(), and checks ValueError when namespaces or nested attributes are missing. Includes helper _build_model.
Changelog CHANGELOG.md	Added Unreleased entry noting fix: State.assign now copies namespaced extended and custom state attributes.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 25.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title directly and accurately describes the main fix: correcting State.assign to properly handle namespaced extended and custom attributes, which is the core objective of this pull request.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

🧹 Nitpick comments (2)

newton/_src/sim/state.py (1)

183-209: LGTM — namespaced extended-attr branch mirrors the top-level semantics correctly.

Branch analysis:

• Both namespaces missing → skip ✓
• One namespace missing while the other holds a wp.array → val_* becomes None via the conditional getattr, and the subsequent val_self is None or not array_self / val_other is None or not array_other checks raise with the qualified "ns.attr" name ✓
• Both namespaces present, neither has the attribute allocated → val_self/val_other both None, skipped by not array_self and not array_other ✓
• Both present and allocated → val_self.assign(val_other) ✓

Optional nit: the top-level loop and this block share the same presence/mismatch/assign pattern. If a third call site ever appears (or recursion is introduced, as noted in the PR description), extracting a small _assign_array(name, val_self, val_other) helper would DRY this up. Not required for this PR.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@newton/_src/sim/state.py` around lines 183 - 209, The namespaced
extended-attribute block duplicates the presence/mismatch/assign pattern used at
top-level; optionally extract a small helper (e.g.,
_assign_array(qualified_name, val_self, val_other)) to encapsulate the checks
and assignment so both the top-level loop and the EXTENDED_ATTRIBUTES loop call
the same logic (replace the repeated checks that compute array_self/array_other,
raise the two ValueError messages with the qualified name, and call
val_self.assign(val_other)).

newton/tests/test_state_assign.py (1)

35-64: Tests cover the three main branches — nice.

A couple of optional additions that would harden regression coverage without much cost:

1. Assert the raised ValueError message contains the qualified name (e.g. "mujoco.qfrc_actuator") so a future refactor that drops the namespace prefix from the error message is caught:

   with self.assertRaisesRegex(ValueError, r"mujoco\.qfrc_actuator"):
       state_0.assign(state_1)

2. A positive no-op test where both states have the mujoco namespace but qfrc_actuator was never requested (i.e., the attribute is None on both sides) would lock in the “skip silently” branch at lines 199-200.

Both are nice-to-have, not blocking.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@newton/tests/test_state_assign.py` around lines 35 - 64, Update
TestStateAssignNamespacedAttributes to tighten assertions and add a no-op branch
test: for the two tests that expect ValueError when a namespaced attribute is
missing, replace assertRaises with self.assertRaisesRegex(ValueError,
r"mujoco\.qfrc_actuator") to ensure the error message mentions the qualified
name when calling state_0.assign(state_1); and add a new test that constructs
state_0 and state_1 where both have a mujoco namespace but qfrc_actuator is None
on both sides, then call state_0.assign(state_1) and assert no exception (a
successful no-op) to cover the "skip silently" branch in the assign
implementation.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Nitpick comments:
In `@newton/_src/sim/state.py`:
- Around line 183-209: The namespaced extended-attribute block duplicates the
presence/mismatch/assign pattern used at top-level; optionally extract a small
helper (e.g., _assign_array(qualified_name, val_self, val_other)) to encapsulate
the checks and assignment so both the top-level loop and the EXTENDED_ATTRIBUTES
loop call the same logic (replace the repeated checks that compute
array_self/array_other, raise the two ValueError messages with the qualified
name, and call val_self.assign(val_other)).

In `@newton/tests/test_state_assign.py`:
- Around line 35-64: Update TestStateAssignNamespacedAttributes to tighten
assertions and add a no-op branch test: for the two tests that expect ValueError
when a namespaced attribute is missing, replace assertRaises with
self.assertRaisesRegex(ValueError, r"mujoco\.qfrc_actuator") to ensure the error
message mentions the qualified name when calling state_0.assign(state_1); and
add a new test that constructs state_0 and state_1 where both have a mujoco
namespace but qfrc_actuator is None on both sides, then call
state_0.assign(state_1) and assert no exception (a successful no-op) to cover
the "skip silently" branch in the assign implementation.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Pro

Run ID: 5ec6a0c5-40c9-467e-aa81-f82f0466fc2f

📥 Commits

Reviewing files that changed from the base of the PR and between 37ad83a and 777d9c2.

📒 Files selected for processing (3)

• CHANGELOG.md
• newton/_src/sim/state.py
• newton/tests/test_state_assign.py

[camevor]

Thanks for spotting and fixing this, @jsw7460! Since the namespace issue also affects custom attributes (to an even larger extent), it would make sense to handle those in the same way. Discovering AttributeNamespace entries on the state object would cover both cases in one pass.

[codecov]

Codecov Report

❌ Patch coverage is 96.77419% with 1 line in your changes missing coverage. Please review.
✅ All tests successful. No failed tests found.

Files with missing lines	Patch %	Lines
newton/_src/sim/state.py	96.77%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

[camevor]

Thanks for updating, this looks good to me!