Import upstream snapshot c19f713c415a699a79d71cd96aa13c3104a05047

Upstream: https://github.com/michaelgillett/mjlab
Upstream-Commit: c19f713c415a699a79d71cd96aa13c3104a05047
Upstream-Branch: main

tests/test_sim_data.py

@@ -0,0 +1,100 @@
+"""Tests for sim data bridge."""
+
+from dataclasses import dataclass
+
+import pytest
+import torch
+import warp as wp
+from conftest import get_test_device
+
+from mjlab.sim.sim_data import TorchArray, WarpBridge
+
+
+@dataclass
+class MockData:
+  arr: wp.array
+  val: float = 1.0
+
+
+@pytest.fixture
+def device():
+  """Test device fixture."""
+  return get_test_device()
+
+
+@pytest.fixture
+def mock_data(device):
+  with wp.ScopedDevice(device):
+    return MockData(arr=wp.array([[1.0, 2.0], [3.0, 4.0]], dtype=wp.float32))
+
+
+def test_torch_array_shares_memory(device):
+  """TorchArray modifications affect underlying warp array."""
+  with wp.ScopedDevice(device):
+    wp_arr = wp.array([1.0, 2.0], dtype=wp.float32)
+  torch_arr = TorchArray(wp_arr)
+  torch_arr[0] = 99.0
+  assert wp_arr.numpy()[0] == 99.0
+
+
+def test_torch_array_ops_work(device):
+  """TorchArray supports PyTorch operations."""
+  with wp.ScopedDevice(device):
+    wp_arr = wp.array([1.0, 2.0], dtype=wp.float32)
+  torch_arr = TorchArray(wp_arr)
+  assert torch.allclose(torch_arr * 2, torch.tensor([2.0, 4.0], device=device))
+  assert torch.sum(torch_arr).item() == 3.0  # type: ignore
+
+
+def test_torch_array_expands_batched_singleton_world_dim(device):
+  """Size-1 world dim expands to nworld when mjwarp marks the array batched."""
+  with wp.ScopedDevice(device):
+    wp_arr = wp.array([[1.0, 2.0]], dtype=wp.float32)
+  setattr(wp_arr, "_is_batched", True)  # noqa: B010
+  torch_arr = TorchArray(wp_arr, nworld=4)
+  assert torch_arr.shape == (4, 2)
+
+
+def test_torch_array_keeps_unbatched_singleton_dim(device):
+  """A real size-1 leading dim without the batched marker is not expanded."""
+  with wp.ScopedDevice(device):
+    wp_arr = wp.array([[1.0, 2.0]], dtype=wp.float32)
+  torch_arr = TorchArray(wp_arr, nworld=4)
+  assert torch_arr.shape == (1, 2)
+
+
+def test_bridge_wraps_arrays(mock_data):
+  """WarpBridge wraps warp arrays as TorchArray."""
+  bridge = WarpBridge(mock_data)
+  assert isinstance(bridge.arr, TorchArray)
+  assert bridge.val == 1.0  # Non-arrays pass through.
+
+
+def test_bridge_preserves_memory_on_slice_assign(mock_data, device):
+  """Slice assignment preserves memory addresses (CUDA graph safe)."""
+  bridge = WarpBridge(mock_data)
+  initial_ptr = bridge.arr._tensor.data_ptr()
+
+  bridge.arr[:] = torch.zeros((2, 2), device=device)
+
+  assert bridge.arr._tensor.data_ptr() == initial_ptr
+  assert torch.all(bridge.arr._tensor == 0)
+
+
+def test_bridge_raises_on_setattr(mock_data, device):
+  """Direct assignment raises AttributeError with helpful message."""
+  bridge = WarpBridge(mock_data)
+
+  with pytest.raises(AttributeError, match="Cannot set attribute 'arr' on WarpBridge"):
+    bridge.arr = torch.zeros((2, 2), device=device)
+
+  with pytest.raises(AttributeError, match="Use in-place operations instead"):
+    bridge.val = 42.0
+
+
+def test_bridge_caches_wrappers(mock_data):
+  """WarpBridge caches wrapper objects."""
+  bridge = WarpBridge(mock_data)
+  arr1 = bridge.arr
+  arr2 = bridge.arr
+  assert arr1 is arr2  # Same cached object.