fix(sim): encode fused root angular reset in root frame

src/unilab/base/reset_state.py

@@ -153,6 +153,7 @@ def __init__(
         self._tensor_root_columns: (
             tuple[tuple[int, ...], tuple[int, ...], torch.Tensor, torch.Tensor] | None
         ) = None
+        self._tensor_root_velocity: torch.Tensor | None = None
         self._tensor_joint_columns: (
             tuple[tuple[int, ...], tuple[int, ...], torch.Tensor, torch.Tensor] | None
         ) = None
@@ -324,6 +325,7 @@ def begin_tensor(
         self._requesting_terms.clear()
         self._last_commit_had_writes = False
         self._tensor_active = True
+        self._ensure_tensor_root_velocity_scratch()
 
     def write_entity_state(
         self,
@@ -1343,7 +1345,7 @@ def write_root_state_tensor(
             qpos_columns=qpos_columns_t,
             qpos_values=root_state[:, :7].contiguous(),
             qvel_columns=qvel_columns_t,
-            qvel_values=root_state[:, 7:13].contiguous(),
+            qvel_values=self._encode_tensor_root_velocity(root_state),
         )
         self._tensor_has_writes = True
 
@@ -1465,7 +1467,7 @@ def write_motion_state_tensor(
             qpos_columns=root_qpos_columns,
             qpos_values=root_state[:, :7].contiguous(),
             qvel_columns=root_qvel_columns,
-            qvel_values=root_state[:, 7:13].contiguous(),
+            qvel_values=self._encode_tensor_root_velocity(root_state),
         )
         transaction.write_state_mask(
             effective_mask,
@@ -1864,6 +1866,37 @@ def _validate_tensor_root_payload(self, root_state: torch.Tensor, *, term_name:
             )
         self._validate_tensor_root_quaternions(root_state[:, 3:7], term_name=term_name)
 
+    def _ensure_tensor_root_velocity_scratch(self) -> None:
+        """Keep the encoded root-velocity carrier at one stable device address."""
+        device = self._packed_reset_device
+        assert device is not None
+        scratch = self._tensor_root_velocity
+        if (
+            scratch is not None
+            and scratch.shape == (self._num_envs, 6)
+            and scratch.dtype is torch.float32
+            and scratch.device == device
+        ):
+            return
+        self._tensor_root_velocity = torch.empty(
+            (self._num_envs, 6), dtype=torch.float32, device=device
+        )
+
+    def _encode_tensor_root_velocity(self, root_state: torch.Tensor) -> torch.Tensor:
+        """Encode community world-frame velocity for generalized free-root qvel.
+
+        Free-root linear velocity is world-frame already. Angular velocity is
+        public world-frame input but generalized ``qvel`` stores it in the root
+        body frame. The full-width carrier avoids masking this conversion to a
+        selected-row allocation while preserving one masked state write.
+        """
+        self._ensure_tensor_root_velocity_scratch()
+        velocity = self._tensor_root_velocity
+        assert velocity is not None
+        velocity[:, :3] = root_state[:, 7:10]
+        velocity[:, 3:6] = quat_apply_inverse(root_state[:, 3:7], root_state[:, 10:13])
+        return velocity.contiguous()
+
     def _validate_tensor_root_quaternions(self, values: torch.Tensor, *, term_name: str) -> None:
         norms = torch.linalg.vector_norm(values, dim=-1)
         if not bool(torch.isclose(norms, torch.ones_like(norms), rtol=1e-5, atol=1e-6).all()):

tests/base/test_backend_conformance.py

@@ -21,6 +21,7 @@
 
 from unilab.assets import ASSETS_ROOT_PATH
 from unilab.base.backend_factory import create_backend
+from unilab.base.reset_state import ResetStateTransaction
 from unilab.base.scene import SceneCfg
 
 REPO_ROOT = Path(__file__).resolve().parents[2]
@@ -436,6 +437,51 @@ def test_root_qvel_body_angular_contract_reads_back_world_velocity(backend_type:
     np.testing.assert_allclose(angular_velocity[0], [1.0, 0.0, 0.0], atol=2e-6)
 
 
+@pytest.mark.parametrize("backend_type", ["mujoco", "motrix"])
+def test_fused_tensor_root_reset_reads_back_world_angular_velocity(
+    backend_type: str,
+) -> None:
+    _require_backend(backend_type)
+    backend = _create_backend_for_test(
+        backend_type,
+        SceneCfg(model_file=_G1_SCENE),
+        1,
+        SIM_DT,
+        base_name="pelvis",
+        add_body_sensors=True,
+    )
+    backend.materialize()
+    try:
+        layout = backend.get_root_state_layout("pelvis")
+        qvel_indices = np.asarray(layout.qvel_indices)
+        reset = ResetStateTransaction(backend)
+        reset.declare_packed_reset_device(torch.device("cpu"))
+        root_state = torch.zeros((1, 13), dtype=torch.float32)
+        half_sqrt = float(np.sqrt(0.5))
+        root_state[:, 3:7] = torch.tensor([half_sqrt, 0.0, 0.0, half_sqrt])
+        root_state[:, 10:13] = torch.tensor([[1.0, 0.0, 0.0]])
+        reset_mask = torch.ones(1, dtype=torch.bool)
+
+        with reset.scoped_device_event_tensor(reset_mask):
+            reset.write_root_state_tensor(
+                reset_mask,
+                layout,
+                root_state,
+                term_name="fused_root_reset_contract",
+            )
+
+        generalized_qvel = backend.get_state_views(("qvel",))["qvel"].detach().cpu().numpy()
+        np.testing.assert_allclose(
+            generalized_qvel[0, qvel_indices[3:6]],
+            [0.0, -1.0, 0.0],
+            atol=2e-6,
+        )
+        world_angular = backend.get_sensor_view("track_angvel_w_pelvis").detach().cpu().numpy()
+        np.testing.assert_allclose(world_angular[0], [1.0, 0.0, 0.0], atol=2e-6)
+    finally:
+        backend.close()
+
+
 def test_mujoco_root_layout_resolves_a_nonfirst_free_joint(tmp_path: Path) -> None:
     pytest.importorskip(
         "unisim.backend.mujoco.backend",

tests/base/test_reset_state.py

@@ -775,6 +775,149 @@ def test_packed_root_state_payloads_remain_contiguous() -> None:
     np.testing.assert_allclose(backend.qvel[[0, 2]], original_qvel[[0, 2]])
 
 
+def test_packed_root_state_encodes_world_angular_velocity_in_root_frame() -> None:
+    backend = _RecordingPackedTensorResetBackend()
+    transaction = _transaction(backend)
+    transaction.declare_packed_reset_device(torch.device("cpu"))
+    reset_mask = torch.tensor([False, True, False, True])
+    original_qvel = backend.qvel.copy()
+    layout = RootStateLayout(tuple(range(7)), tuple(range(6)))
+    root_state = torch.zeros((4, 13), dtype=torch.float32)
+    quaternion = torch.tensor(
+        [torch.sqrt(torch.tensor(0.5)), 0.0, 0.0, torch.sqrt(torch.tensor(0.5))],
+        dtype=torch.float32,
+    )
+    root_state[:, 3:7] = quaternion
+    root_state[1, 7:10] = torch.tensor([1.0, 2.0, 3.0])
+    root_state[3, 7:10] = torch.tensor([4.0, 5.0, 6.0])
+    root_state[1, 10:13] = torch.tensor([1.0, 0.0, 0.0])
+    root_state[3, 10:13] = torch.tensor([0.0, 0.0, 2.0])
+
+    with transaction.scoped_device_event_tensor(reset_mask):
+        transaction.write_root_state_tensor(reset_mask, layout, root_state, term_name="root_reset")
+
+    assert len(backend.state_write_payloads) == 1
+    _, encoded_velocity = backend.state_write_payloads[0]
+    torch.testing.assert_close(encoded_velocity[:, :3], root_state[[1, 3], 7:10])
+    torch.testing.assert_close(
+        encoded_velocity[:, 3:6],
+        torch.tensor([[0.0, -1.0, 0.0], [0.0, 0.0, 2.0]]),
+        rtol=0,
+        atol=2e-7,
+    )
+    _, _, committed_qvel, _ = backend.tensor_reset_calls[0]
+    torch.testing.assert_close(committed_qvel[:, :6], encoded_velocity, rtol=0, atol=0)
+    np.testing.assert_allclose(backend.qvel[[0, 2]], original_qvel[[0, 2]], rtol=0, atol=1e-7)
+
+
+def test_packed_motion_tensor_encodes_world_angular_velocity_in_root_frame() -> None:
+    backend = _RecordingPackedTensorResetBackend()
+    transaction = _transaction(backend)
+    transaction.declare_packed_reset_device(torch.device("cpu"))
+    reset_mask = torch.tensor([True, False, True, False])
+    layout = RootStateLayout(tuple(range(7)), tuple(range(6)))
+    root_state = torch.zeros((4, 13), dtype=torch.float32)
+    quaternion = torch.tensor(
+        [torch.sqrt(torch.tensor(0.5)), 0.0, 0.0, torch.sqrt(torch.tensor(0.5))],
+        dtype=torch.float32,
+    )
+    root_state[:, 3:7] = quaternion
+    root_state[[0, 2], 7:10] = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
+    root_state[[0, 2], 10:13] = torch.tensor([[1.0, 0.0, 0.0], [0.0, 0.0, 2.0]])
+    joint_position = torch.zeros((4, 2), dtype=torch.float32)
+    joint_velocity = torch.zeros((4, 2), dtype=torch.float32)
+
+    with transaction.scoped_device_event_tensor(reset_mask):
+        transaction.write_motion_state_tensor(
+            reset_mask,
+            layout,
+            np.array([7, 8], dtype=np.int64),
+            np.array([6, 7], dtype=np.int64),
+            root_state,
+            joint_position,
+            joint_velocity,
+            term_name="motion_reset",
+        )
+
+    assert len(backend.state_write_payloads) == 2
+    _, encoded_velocity = backend.state_write_payloads[0]
+    torch.testing.assert_close(encoded_velocity[:, :3], root_state[[0, 2], 7:10])
+    torch.testing.assert_close(
+        encoded_velocity[:, 3:6],
+        torch.tensor([[0.0, -1.0, 0.0], [0.0, 0.0, 2.0]]),
+        rtol=0,
+        atol=2e-7,
+    )
+    assert len(backend.tensor_reset_calls) == 1
+
+
+def test_packed_tensor_reset_stays_off_host_and_reuses_velocity_scratch(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    backend = _RecordingPackedTensorResetBackend()
+    transaction = _transaction(backend)
+    transaction.declare_packed_reset_device(torch.device("cpu"))
+    layout = RootStateLayout(tuple(range(7)), tuple(range(6)))
+    root_state = torch.zeros((4, 13), dtype=torch.float32)
+    quaternion = torch.tensor(
+        [torch.sqrt(torch.tensor(0.5)), 0.0, 0.0, torch.sqrt(torch.tensor(0.5))],
+        dtype=torch.float32,
+    )
+    root_state[:, 3:7] = quaternion
+    root_state[:, 10] = 1.0
+
+    def fail_host_call(*_args, **_kwargs):
+        raise AssertionError("fused tensor reset must not cross a Torch host boundary")
+
+    with (
+        transaction.scoped_device_event_tensor(torch.ones(4, dtype=torch.bool)),
+        monkeypatch.context() as patches,
+    ):
+        patches.setattr(torch.Tensor, "cpu", fail_host_call)
+        patches.setattr(torch.Tensor, "numpy", fail_host_call)
+        patches.setattr(torch.Tensor, "item", fail_host_call)
+        patches.setattr(torch.Tensor, "tolist", fail_host_call)
+        transaction.write_root_state_tensor(
+            torch.ones(4, dtype=torch.bool), layout, root_state, term_name="root_reset"
+        )
+        scratch = transaction._tensor_root_velocity
+        transaction.write_root_state_tensor(
+            torch.ones(4, dtype=torch.bool), layout, root_state, term_name="root_reset_2"
+        )
+        assert transaction._tensor_root_velocity is scratch
+
+    assert len(backend.tensor_reset_calls) == 1
+
+
+def test_packed_motion_tensor_reset_preserves_owner_subset_write_boundary() -> None:
+    backend = _RecordingPackedTensorResetBackend()
+    transaction = _transaction(backend)
+    transaction.declare_packed_reset_device(torch.device("cpu"))
+    active_mask = torch.tensor([False, True, True, False])
+    owner_mask = torch.tensor([False, False, True, False])
+    layout = RootStateLayout(tuple(range(7)), tuple(range(6)))
+    root_state = torch.zeros((4, 13), dtype=torch.float32)
+    root_state[:, 3] = 1.0
+    joint_position = torch.zeros((4, 2), dtype=torch.float32)
+    joint_velocity = torch.zeros((4, 2), dtype=torch.float32)
+
+    with transaction.scoped_device_owner_tensor(active_mask):
+        transaction.write_motion_state_tensor(
+            owner_mask,
+            layout,
+            np.array([7, 8], dtype=np.int64),
+            np.array([6, 7], dtype=np.int64),
+            root_state,
+            joint_position,
+            joint_velocity,
+            term_name="owner_motion_reset",
+        )
+
+    assert backend.state_write_payloads == []
+    assert backend.tensor_reset_calls == []
+    assert not transaction.last_commit_had_writes
+
+
 @pytest.mark.parametrize(
     ("root_state", "error", "match"),
     [