diff --git a/src/unilab/base/reset_state.py b/src/unilab/base/reset_state.py index 3d78890d1..d472a98b4 100644 --- a/src/unilab/base/reset_state.py +++ b/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()): diff --git a/tests/base/test_backend_conformance.py b/tests/base/test_backend_conformance.py index 4b11172c0..70bc114b9 100644 --- a/tests/base/test_backend_conformance.py +++ b/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", diff --git a/tests/base/test_reset_state.py b/tests/base/test_reset_state.py index 6923331b2..ece6c7a0c 100644 --- a/tests/base/test_reset_state.py +++ b/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"), [