Fix root shapes in ArticulationView with fixed base (#1639)

Signed-off-by: Eric Heiden <eric-heiden@outlook.com>
Co-authored-by: Eric Heiden <eric-heiden@outlook.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: 3 people

newton/_src/utils/selection.py

@@ -15,6 +15,7 @@
 
 import functools
 from fnmatch import fnmatch
+from types import NoneType
 from typing import Any
 
 import warp as wp
@@ -1004,38 +1005,29 @@ def _get_attribute_array(self, name: str, source: Model | State | Control, _slic
         # handle custom slice
         if isinstance(_slice, Slice):
             _slice = _slice.get()
-        elif isinstance(_slice, int):
-            _slice = slice(_slice, _slice + 1)
+        elif not isinstance(_slice, (NoneType, int, slice)):
+            raise ValueError(f"Invalid slice type: expected slice or int, got {type(_slice)}")
 
         if _slice is None:
+            value_slice = layout.indices if is_indexed else layout.slice
             value_count = layout.value_count
-            if is_indexed:
-                value_slice = layout.indices
-            else:
-                value_slice = layout.slice
         else:
-            value_count = _slice.stop - _slice.start
-            if is_indexed:
-                value_slice = layout.indices[_slice]
-            else:
-                value_slice = _slice
-
-        shape = (self.world_count, self.count_per_world, value_count)
-        strides = (
-            layout.stride_between_worlds * value_stride,
-            layout.stride_within_worlds * value_stride,
-            value_stride,
-        )
-        slices = (slice(self.world_count), slice(self.count_per_world), value_slice)
+            value_slice = _slice
+            value_count = 1 if isinstance(_slice, int) else _slice.stop - _slice.start
 
         # trailing dimensions for multidimensional attributes
         trailing_shape = attrib.shape[1:]
         trailing_strides = attrib.strides[1:]
         trailing_slices = [slice(s) for s in trailing_shape]
 
-        shape = (*shape, *trailing_shape)
-        strides = (*strides, *trailing_strides)
-        slices = (*slices, *trailing_slices)
+        shape = (self.world_count, self.count_per_world, value_count, *trailing_shape)
+        strides = (
+            layout.stride_between_worlds * value_stride,
+            layout.stride_within_worlds * value_stride,
+            value_stride,
+            *trailing_strides,
+        )
+        slices = (slice(self.world_count), slice(self.count_per_world), value_slice, *trailing_slices)
 
         # construct reshaped attribute array
         attrib = wp.array(

newton/tests/test_selection.py

@@ -47,20 +47,26 @@ def test_empty_selection(self):
         self.assertEqual(selection.get_dof_velocities(model).shape, (1, 1, 0))
         self.assertEqual(selection.get_dof_forces(control).shape, (1, 1, 0))
 
-    def test_selection_shapes(self):
+    def _test_selection_shapes(self, floating: bool):
         # load articulation
         ant = newton.ModelBuilder()
         ant.add_mjcf(
             newton.examples.get_asset("nv_ant.xml"),
             ignore_names=["floor", "ground"],
+            floating=floating,
         )
 
         L = 9  # num links
         J = 9  # num joints
-        D = 14  # num joint dofs
-        C = 15  # num joint coords
         S = 13  # num shapes
 
+        if floating:
+            D = 14  # num joint dofs
+            C = 15  # num joint coords
+        else:
+            D = 8  # num joint dofs
+            C = 8  # num joint coords
+
         # scene with just one ant
         single_ant_model = ant.finalize()
 
@@ -69,7 +75,10 @@ def test_selection_shapes(self):
         self.assertEqual(single_ant_view.world_count, 1)
         self.assertEqual(single_ant_view.count_per_world, 1)
         self.assertEqual(single_ant_view.get_root_transforms(single_ant_model).shape, (1, 1))
-        self.assertEqual(single_ant_view.get_root_velocities(single_ant_model).shape, (1, 1))
+        if floating:
+            self.assertEqual(single_ant_view.get_root_velocities(single_ant_model).shape, (1, 1))
+        else:
+            self.assertIsNone(single_ant_view.get_root_velocities(single_ant_model))
         self.assertEqual(single_ant_view.get_link_transforms(single_ant_model).shape, (1, 1, L))
         self.assertEqual(single_ant_view.get_link_velocities(single_ant_model).shape, (1, 1, L))
         self.assertEqual(single_ant_view.get_dof_positions(single_ant_model).shape, (1, 1, C))
@@ -92,7 +101,10 @@ def test_selection_shapes(self):
         self.assertEqual(single_ant_per_world_view.world_count, W)
         self.assertEqual(single_ant_per_world_view.count_per_world, 1)
         self.assertEqual(single_ant_per_world_view.get_root_transforms(single_ant_per_world_model).shape, (W, 1))
-        self.assertEqual(single_ant_per_world_view.get_root_velocities(single_ant_per_world_model).shape, (W, 1))
+        if floating:
+            self.assertEqual(single_ant_per_world_view.get_root_velocities(single_ant_per_world_model).shape, (W, 1))
+        else:
+            self.assertIsNone(single_ant_per_world_view.get_root_velocities(single_ant_per_world_model))
         self.assertEqual(single_ant_per_world_view.get_link_transforms(single_ant_per_world_model).shape, (W, 1, L))
         self.assertEqual(single_ant_per_world_view.get_link_velocities(single_ant_per_world_model).shape, (W, 1, L))
         self.assertEqual(single_ant_per_world_view.get_dof_positions(single_ant_per_world_model).shape, (W, 1, C))
@@ -128,7 +140,10 @@ def test_selection_shapes(self):
         self.assertEqual(multi_ant_per_world_view.world_count, W)
         self.assertEqual(multi_ant_per_world_view.count_per_world, A)
         self.assertEqual(multi_ant_per_world_view.get_root_transforms(multi_ant_per_world_model).shape, (W, A))
-        self.assertEqual(multi_ant_per_world_view.get_root_velocities(multi_ant_per_world_model).shape, (W, A))
+        if floating:
+            self.assertEqual(multi_ant_per_world_view.get_root_velocities(multi_ant_per_world_model).shape, (W, A))
+        else:
+            self.assertIsNone(multi_ant_per_world_view.get_root_velocities(multi_ant_per_world_model))
         self.assertEqual(multi_ant_per_world_view.get_link_transforms(multi_ant_per_world_model).shape, (W, A, L))
         self.assertEqual(multi_ant_per_world_view.get_link_velocities(multi_ant_per_world_model).shape, (W, A, L))
         self.assertEqual(multi_ant_per_world_view.get_dof_positions(multi_ant_per_world_model).shape, (W, A, C))
@@ -149,6 +164,12 @@ def test_selection_shapes(self):
             multi_ant_per_world_view.get_attribute("shape_thickness", multi_ant_per_world_model).shape, (W, A, S)
         )
 
+    def test_selection_shapes_floating_base(self):
+        self._test_selection_shapes(floating=True)
+
+    def test_selection_shapes_fixed_base(self):
+        self._test_selection_shapes(floating=False)
+
     def test_selection_shape_values_noncontiguous(self):
         """Test that shape attribute values are correct when shape selection is non-contiguous."""
         # Build a 3-link chain: base -> link1 -> link2