add_gaussians.py

    def add_gaussians(self, env_idx: int):
        if "gs_renderer" not in self.env.sensors:
            return
        ply_renderers = self.env.sensors["gs_renderer"].get_gs_renderers()
        mesh_name = self.scene_manager.mesh_names[self.scene_manager.mesh_id_for_env_id(env_idx)]
        splat_name = '/'.join(mesh_name.split('/')[:-1])
        renderer = ply_renderers[splat_name]
        centers = renderer.means.cpu().numpy() / renderer.dataparser_scale
        colors = renderer.colors_viser.cpu().numpy()
        opacities = renderer.opacities.cpu().numpy()[..., None]
        scales = renderer.scales.cpu().numpy() / renderer.dataparser_scale
        Rs = vtf.SO3(renderer.quats.cpu().numpy()).as_matrix()
        covariances = np.einsum(
            "nij,njk,nlk->nil", Rs, np.eye(3)[None, :, :] * scales[:, None, :] ** 2, Rs
        )

        cam_offset = self.scene_manager.cam_offset[env_idx].cpu().numpy()
        env_offset = self.env.env_origins[env_idx].cpu().numpy()

        ig_to_orig_rot = vtf.SO3.from_matrix(self.scene_manager.ig_to_orig_rot[env_idx].cpu().numpy())
        dataparser_transform = vtf.SE3.from_matrix(renderer.dataparser_transform.cpu().numpy())

        # Move from GS frame to IG frame.
        cumulative_transform = dataparser_transform.inverse()
        cumulative_transform = vtf.SE3.from_rotation_and_translation(
            vtf.SO3.identity(),
            cam_offset
        ).inverse().multiply(cumulative_transform)
        cumulative_transform = vtf.SE3.from_rotation_and_translation(
            ig_to_orig_rot,
            np.zeros(3)
        ).inverse().multiply(cumulative_transform)
        cumulative_transform = vtf.SE3.from_rotation_and_translation(
              vtf.SO3.identity(),
              env_offset
        ).multiply(cumulative_transform)

        if self._gs_handle is not None:
            self._gs_handle.remove()
        self._gs_handle = self.server.scene.add_gaussian_splats(
            "/gs",
            centers=centers,
            covariances=covariances,
            rgbs=colors,
            opacities=opacities,
            wxyz=cumulative_transform.rotation().wxyz,
            position=cumulative_transform.translation(),
        )