Render Gaussian splat shapes in ViewerViser (#4018)

Author: aneangel

changelog/2099.fixed.md

@@ -0,0 +1 @@
+Render `newton.Gaussian` splat shapes in `ViewerViser`. `log_gaussian()` was previously a no-op there, so Gaussian splat assets silently vanished when viewing with `--viewer viser`. Local-space centers, covariances, colors, and opacities are uploaded to viser's native Gaussian splat renderer once per asset and cached; only the node's position and orientation are updated on subsequent frames.

newton/_src/viewer/viewer_viser.py

@@ -37,6 +37,7 @@ class ViewerViser(ViewerBase):
     """
 
     _viser_module = None
+    _SH_C0 = 0.28209479177387814
 
     @classmethod
     def _get_viser(cls):
@@ -150,6 +151,7 @@ def __init__(
         self._meshes = {}
         self._instances = {}
         self._scene_handles = {}  # Track viser scene node handles
+        self._gaussian_splats = {}  # Track cached Gaussian splat upload keys, by name
         self._line_segment_counts = {}
         self._line_versions = {}
 
@@ -194,6 +196,16 @@ def clear_model(self):
                 self._remove_plane_handles(plane_name)
         self._plane_meshes = {name: value for name, value in self._plane_meshes.items() if not owns(name)}
 
+        for gaussian_name in list(getattr(self, "_gaussian_splats", {}).keys()):
+            if owns(gaussian_name):
+                handle = self._scene_handles.pop(gaussian_name, None)
+                if handle is not None:
+                    try:
+                        handle.remove()
+                    except Exception:
+                        pass
+                self._gaussian_splats.pop(gaussian_name, None)
+
         for name, handle in list(getattr(self, "_scene_handles", {}).items()):
             if not owns(name):
                 continue
@@ -204,6 +216,7 @@ def clear_model(self):
             self._scene_handles.pop(name, None)
             self._instances.pop(name, None)
             self._meshes.pop(name, None)
+            self._gaussian_splats.pop(name, None)
             self._line_segment_counts.pop(name, None)
             self._line_versions.pop(name, None)
 
@@ -675,6 +688,7 @@ def log_mesh(
                 self._scene_handles[name].remove()
             except Exception:
                 pass
+            del self._scene_handles[name]
 
         if hidden:
             return
@@ -715,6 +729,25 @@ def _quats_xyzw_to_wxyz(quats_xyzw: np.ndarray) -> np.ndarray:
         quats_wxyz[:, 3] = quats_xyzw[:, 2]
         return quats_wxyz[0] if was_1d else quats_wxyz
 
+    @staticmethod
+    def _quats_xyzw_to_rotmats(quats_xyzw: np.ndarray) -> np.ndarray:
+        """Convert a batch of XYZW quaternions to (N, 3, 3) rotation matrices."""
+        quats_xyzw = np.asarray(quats_xyzw, dtype=np.float32)
+        quats_xyzw = quats_xyzw / np.maximum(np.linalg.norm(quats_xyzw, axis=1, keepdims=True), 1e-12)
+        x, y, z, w = quats_xyzw[:, 0], quats_xyzw[:, 1], quats_xyzw[:, 2], quats_xyzw[:, 3]
+        n = quats_xyzw.shape[0]
+        rot = np.empty((n, 3, 3), dtype=np.float32)
+        rot[:, 0, 0] = 1.0 - 2.0 * (y * y + z * z)
+        rot[:, 0, 1] = 2.0 * (x * y - w * z)
+        rot[:, 0, 2] = 2.0 * (x * z + w * y)
+        rot[:, 1, 0] = 2.0 * (x * y + w * z)
+        rot[:, 1, 1] = 1.0 - 2.0 * (x * x + z * z)
+        rot[:, 1, 2] = 2.0 * (y * z - w * x)
+        rot[:, 2, 0] = 2.0 * (x * z - w * y)
+        rot[:, 2, 1] = 2.0 * (y * z + w * x)
+        rot[:, 2, 2] = 1.0 - 2.0 * (x * x + y * y)
+        return rot
+
     def _remove_plane_handles(self, name: str):
         """Remove any plane-grid handles associated with an instance batch."""
         handle = self._plane_handles.pop(name, None)
@@ -1265,6 +1298,7 @@ def log_points(
                 self._scene_handles[name].remove()
             except Exception:
                 pass
+            del self._scene_handles[name]
 
         if hidden:
             return
@@ -1313,6 +1347,93 @@ def log_points(
         )
         self._scene_handles[name] = handle
 
+    @override
+    def log_gaussian(
+        self,
+        name: str,
+        gaussian: newton.Gaussian | None,
+        xform: wp.transformf | None = None,
+        hidden: bool = False,
+    ):
+        """
+        Log a :class:`newton.Gaussian` splat asset using viser's native Gaussian renderer.
+
+        Note: viser's ``add_gaussian_splats`` is marked experimental upstream and its
+        API may change or be removed in a future viser release.
+
+        Args:
+            name: Unique path/name for the Gaussian splat asset.
+            gaussian: The :class:`newton.Gaussian` asset to visualize, with centers and
+                per-axis scales in meters [m]. ``None`` removes any existing asset at ``name``.
+            xform: Optional world-space transform applied to the splat asset; its
+                translation component is in meters [m].
+            hidden: Whether the splat asset should be hidden.
+        """
+        name = self._qualify(name)
+
+        if gaussian is None or gaussian.count == 0:
+            if name in self._scene_handles:
+                try:
+                    self._scene_handles[name].remove()
+                except Exception:
+                    pass
+                self._scene_handles.pop(name, None)
+            self._gaussian_splats.pop(name, None)
+            return
+
+        if hidden:
+            if name in self._scene_handles:
+                self._scene_handles[name].visible = False
+            return
+
+        cached = self._gaussian_splats.get(name)
+
+        if (
+            cached is None
+            or cached["gaussian"] is not gaussian
+            or cached["count"] != gaussian.count
+            or self._scene_handles.get(name) is not cached["handle"]
+        ):
+            centers = self._to_numpy(gaussian.positions).astype(np.float32)
+            rotations_xyzw = self._to_numpy(gaussian.rotations).astype(np.float32)
+            scales = self._to_numpy(gaussian.scales).astype(np.float32)
+            opacities = self._to_numpy(gaussian.opacities).astype(np.float32).reshape(-1, 1)
+
+            # Local-space covariance per Gaussian: Sigma = R * diag(scale^2) * R^T
+            rot_mats = self._quats_xyzw_to_rotmats(rotations_xyzw)
+            scale_sq = scales * scales
+            covariances = np.einsum("nij,nj,nkj->nik", rot_mats, scale_sq, rot_mats).astype(np.float32)
+
+            sh_coeffs = self._to_numpy(gaussian.sh_coeffs)
+            if sh_coeffs is not None and sh_coeffs.shape[1] >= 3:
+                rgbs = np.clip(self._SH_C0 * sh_coeffs[:, :3] + 0.5, 0.0, 1.0).astype(np.float32)
+            else:
+                rgbs = np.ones((gaussian.count, 3), dtype=np.float32)
+
+            if cached is not None and name in self._scene_handles:
+                try:
+                    self._scene_handles[name].remove()
+                except Exception:
+                    pass
+
+            handle = self._call_scene_method(
+                self._server.scene.add_gaussian_splats,
+                name=name,
+                centers=centers,
+                covariances=covariances,
+                rgbs=rgbs,
+                opacities=opacities,
+            )
+            self._scene_handles[name] = handle
+            self._gaussian_splats[name] = {"gaussian": gaussian, "count": gaussian.count, "handle": handle}
+
+        handle = self._scene_handles[name]
+        handle.visible = True
+        if xform is not None:
+            xform_np = np.asarray(xform, dtype=np.float32)
+            handle.position = xform_np[:3]
+            handle.wxyz = self._quats_xyzw_to_wxyz(xform_np[3:7])
+
     @override
     def log_array(self, name: str, array: wp.array[Any] | np.ndarray):
         """Viser viewer does not visualize generic arrays.