Add com points in viewer

newton/_src/viewer/kernels.py

@@ -427,6 +427,23 @@ def compute_joint_basis_lines(
     line_colors[tid] = color
 
 
+@wp.kernel
+def compute_com_positions(
+    body_q: wp.array(dtype=wp.transform),
+    body_com: wp.array(dtype=wp.vec3),
+    body_world: wp.array(dtype=int),
+    world_offsets: wp.array(dtype=wp.vec3),
+    com_positions: wp.array(dtype=wp.vec3),
+):
+    tid = wp.tid()
+    body_tf = body_q[tid]
+    world_com = wp.transform_point(body_tf, body_com[tid])
+    world_idx = body_world[tid]
+    if world_offsets and world_idx >= 0 and world_idx < world_offsets.shape[0]:
+        world_com = world_com + world_offsets[world_idx]
+    com_positions[tid] = world_com
+
+
 @wp.func
 def depth_to_color(depth: float, min_depth: float, max_depth: float) -> wp.vec3:
     """Convert depth value to a color using a blue-to-red colormap."""

newton/_src/viewer/viewer.py

@@ -63,6 +63,10 @@ def __init__(self):
         self._joint_points1 = None
         self._joint_colors = None
 
+        self._com_positions = None
+        self._com_colors = None
+        self._com_radii = None
+
         # World offset support
         self.world_offsets = None  # Array of vec3 offsets per world
 
@@ -342,6 +346,7 @@ def log_state(self, state):
         self._log_triangles(state)
         self._log_particles(state)
         self._log_joints(state)
+        self._log_com(state)
 
         self.model_changed = False
 
@@ -1212,6 +1217,33 @@ def _log_joints(self, state):
         # Log all joint lines in a single call
         self.log_lines("/model/joints", self._joint_points0, self._joint_points1, self._joint_colors)
 
+    def _log_com(self, state):
+        num_bodies = self.model.body_count
+        if num_bodies == 0:
+            return
+
+        if self._com_positions is None or len(self._com_positions) < num_bodies:
+            self._com_positions = wp.zeros(num_bodies, dtype=wp.vec3, device=self.device)
+            self._com_colors = wp.full(num_bodies, wp.vec3(1.0, 0.8, 0.0), device=self.device)
+            self._com_radii = wp.full(num_bodies, 0.05, dtype=float, device=self.device)
+
+        from .kernels import compute_com_positions  # noqa: PLC0415
+
+        wp.launch(
+            kernel=compute_com_positions,
+            dim=num_bodies,
+            inputs=[
+                state.body_q,
+                self.model.body_com,
+                self.model.body_world,
+                self.world_offsets,
+            ],
+            outputs=[self._com_positions],
+            device=self.device,
+        )
+
+        self.log_points("/model/com", self._com_positions, self._com_radii, self._com_colors, hidden=not self.show_com)
+
     def _log_triangles(self, state):
         if self.model.tri_count:
             self.log_mesh(

newton/_src/viewer/viewer_usd.py

@@ -487,6 +487,8 @@ def log_points(self, name, points, radii, colors, hidden=False):
 
         if colors is None:
             color_interp = "constant"
+        elif isinstance(colors, wp.array):
+            color_interp = "vertex"
         elif len(colors) == 3 and all(np.isscalar(x) for x in colors):
             color_interp = "constant"
         else: