Switch to viser

[Gregwar]

Currently, placo visualization uses meshcat

Viser seems more efficient, with more possibilities for user interaction and customization

[qiang-dyna]

"""
Viser-based visualization module for placo robots.

This module provides a similar API to visualization.py but uses viser instead of meshcat.
All functions maintain the same interface for easy migration.

Uses ViserUrdf for automatic URDF loading and visualization, which handles:
- Mesh file loading (STL, OBJ, etc.)
- Primitive geometries (BOX, CYLINDER, SPHERE)
- Scene graph management
- Efficient joint updates
"""

from pathlib import Path

import numpy as np
import viser

import placo

try:
    import trimesh

    HAS_TRIMESH = True
except ImportError:
    HAS_TRIMESH = False
    trimesh = None

# Import ViserUrdf (required)
try:
    from viser.extras import ViserUrdf

    HAS_VISER_URDF = True
except ImportError:
    HAS_VISER_URDF = False
    ViserUrdf = None

server: viser.ViserServer | None = None
robot_names: dict = {}
robot_viz_handles: dict = {}  # Store viser handles for robot meshes
visualization_objects: dict = {}  # Store handles for various visualization objects


def get_viewer() -> viser.ViserServer:
    """
    Gets the viser server, if it doesn't exist, create it
    """
    global server

    if server is None:
        server = viser.ViserServer()
        port = getattr(server, "port", 8080)
        print(f"Viser server started. Open http://localhost:{port} in your web browser")

    return server


class RobotViz:
    """Wrapper to maintain compatibility with pinocchio's display() method using ViserUrdf"""

    def __init__(
        self,
        robot: placo.RobotWrapper,
        name: str,
        server: viser.ViserServer,
        urdf_path: Path | str | None = None,
        load_collision: bool = False,
    ):
        if not HAS_VISER_URDF:
            raise ImportError(
                "ViserUrdf is required but not available. "
                "Please ensure viser.extras.ViserUrdf is importable."
            )

        self.robot = robot
        self.name = name
        self.server = server
        self.viser_urdf: ViserUrdf | None = None

        # Get URDF path - use provided path or try to find it
        if urdf_path is not None:
            urdf_path = Path(urdf_path)
            if not urdf_path.is_absolute():
                urdf_path = urdf_path.resolve()
        else:
            urdf_path = self._find_urdf_path()

        if urdf_path is None or not urdf_path.exists():
            raise FileNotFoundError(
                f"Could not find URDF file for robot. "
                f"Robot model_directory: {getattr(robot, 'model_directory', 'unknown')}. "
                f"Please provide urdf_path parameter to robot_viz()."
            )

        # Create ViserUrdf
        self.viser_urdf = ViserUrdf(
            self.server,
            urdf_or_path=urdf_path,
            load_meshes=True,
            load_collision_meshes=load_collision,
            collision_mesh_color_override=(1.0, 0.0, 0.0, 0.5),
            root_node_name=f"/{self.name}",
        )
        print(f"Loaded robot using ViserUrdf from {urdf_path}")

    def _find_urdf_path(self) -> Path | None:
        """Find the URDF file path from the robot's model directory"""
        # Try to get model_directory from robot
        if hasattr(self.robot, "model_directory") and self.robot.model_directory:
            model_dir_str = self.robot.model_directory
            if model_dir_str and model_dir_str != "unknown":
                urdf_path = Path(model_dir_str)

                # If it's already a .urdf file, return it
                if urdf_path.exists() and urdf_path.suffix == ".urdf":
                    return urdf_path

                # Try to find URDF file in the directory
                if urdf_path.is_dir():
                    urdf_dir = urdf_path
                else:
                    urdf_dir = urdf_path.parent

                # Try common URDF filenames
                urdf_filenames = ["robot.urdf", "model.urdf"]
                # If the model_directory path name looks like a URDF filename, try it first
                if urdf_path.name.endswith(".urdf"):
                    urdf_filenames.insert(0, urdf_path.name)

                for urdf_filename in urdf_filenames:
                    urdf_file = urdf_dir / urdf_filename
                    if urdf_file.exists():
                        return urdf_file

                # Look for any .urdf file in the directory
                urdf_files = list(urdf_dir.glob("*.urdf"))
                if urdf_files:
                    return urdf_files[0]

        return None

    def display(self, q: np.ndarray):
        """Update robot visualization with new joint configuration"""
        # Update robot state
        self.robot.state.q = q
        self.robot.update_kinematics()

        if self.viser_urdf is None:
            return

        # Map placo joints to viser joints
        actuated_joint_names = self.viser_urdf.get_actuated_joint_names()
        viser_config = []

        for joint_name in actuated_joint_names:
            if joint_name in self.robot.joint_names():
                viser_config.append(self.robot.get_joint(joint_name))
            else:
                # Joint not in placo, use default value
                viser_config.append(0.0)

        # Update viser visualization
        self.viser_urdf.update_cfg(np.array(viser_config))


def robot_viz(
    robot: placo.RobotWrapper,
    name: str = "robot",
    urdf_path: Path | str | None = None,
    load_collision: bool = False,
) -> RobotViz:
    """
    Builds a viser-based robot visualizer using ViserUrdf

    The robot can further be displayed using:

    > viz.display(q)

    Args:
        robot: The placo robot wrapper
        name: Name for the robot in the scene
        urdf_path: Optional path to the URDF file. If not provided, will try to
            find it from the robot's model_directory. It's recommended to provide
            this if the robot was loaded from a specific URDF file.
        load_collision: If True, also load collision meshes
    """
    global robot_names, robot_viz_handles

    robot_names[robot] = name
    server = get_viewer()

    viz = RobotViz(robot, name, server, urdf_path=urdf_path, load_collision=load_collision)
    robot_viz_handles[robot] = viz

    return viz


def frame_viz(name: str, T: np.ndarray, opacity: float = 1.0, scale: float = 1.0) -> None:
    """
    Visualizes a given frame using viser's add_frame API
    """
    global visualization_objects
    server = get_viewer()

    # Extract position and rotation from transform matrix
    position = T[:3, 3]
    rotation_matrix = T[:3, :3]

    # Convert rotation matrix to quaternion
    trace = np.trace(rotation_matrix)
    if trace > 0:
        s = np.sqrt(trace + 1.0) * 2
        w = 0.25 * s
        x = (rotation_matrix[2, 1] - rotation_matrix[1, 2]) / s
        y = (rotation_matrix[0, 2] - rotation_matrix[2, 0]) / s
        z = (rotation_matrix[1, 0] - rotation_matrix[0, 1]) / s
    else:
        if (
            rotation_matrix[0, 0] > rotation_matrix[1, 1]
            and rotation_matrix[0, 0] > rotation_matrix[2, 2]
        ):
            s = (
                np.sqrt(1.0 + rotation_matrix[0, 0] - rotation_matrix[1, 1] - rotation_matrix[2, 2])
                * 2
            )
            w = (rotation_matrix[2, 1] - rotation_matrix[1, 2]) / s
            x = 0.25 * s
            y = (rotation_matrix[0, 1] + rotation_matrix[1, 0]) / s
            z = (rotation_matrix[0, 2] + rotation_matrix[2, 0]) / s
        elif rotation_matrix[1, 1] > rotation_matrix[2, 2]:
            s = (
                np.sqrt(1.0 + rotation_matrix[1, 1] - rotation_matrix[0, 0] - rotation_matrix[2, 2])
                * 2
            )
            w = (rotation_matrix[0, 2] - rotation_matrix[2, 0]) / s
            x = (rotation_matrix[0, 1] + rotation_matrix[1, 0]) / s
            y = 0.25 * s
            z = (rotation_matrix[1, 2] + rotation_matrix[2, 1]) / s
        else:
            s = (
                np.sqrt(1.0 + rotation_matrix[2, 2] - rotation_matrix[0, 0] - rotation_matrix[1, 1])
                * 2
            )
            w = (rotation_matrix[1, 0] - rotation_matrix[0, 1]) / s
            x = (rotation_matrix[0, 2] + rotation_matrix[2, 0]) / s
            y = (rotation_matrix[1, 2] + rotation_matrix[2, 1]) / s
            z = 0.25 * s

    quaternion = (w, x, y, z)

    # Use viser's built-in frame visualization
    node_name = f"frames/{name}"

    if node_name not in visualization_objects:
        # Create frame with scaled axes
        handle = server.scene.add_frame(
            node_name,
            show_axes=True,
            axes_length=0.1 * scale,
            axes_radius=0.005 * scale,
            wxyz=quaternion,
            position=position,
            visible=True,
        )
        visualization_objects[node_name] = handle
    else:
        # Update existing frame
        handle = visualization_objects[node_name]
        handle.position = position
        handle.wxyz = quaternion
        # Update scale if needed
        handle.axes_length = 0.1 * scale
        handle.axes_radius = 0.005 * scale
        handle.visible = True  # Ensure it's visible
        handle.show_axes = True  # Ensure axes are shown


def point_viz(
    name: str,
    point: np.ndarray,
    radius: float = 0.01,
    color: float = 0xFF0000,
    opacity: float = 1.0,
) -> None:
    """
    Visualizes a point (sphere)
    """
    server = get_viewer()
    handle_name = f"point/{name}"

    # Convert hex color to RGB tuple
    r = (color >> 16) & 0xFF
    g = (color >> 8) & 0xFF
    b = color & 0xFF
    rgb_color = (r, g, b)

    try:
        if handle_name in visualization_objects:
            handle = visualization_objects[handle_name]
            handle.position = point
            handle.radius = radius
        else:
            handle = server.scene.add_icosphere(
                handle_name,
                radius=radius,
                position=point,
                color=rgb_color,
                opacity=opacity if opacity < 1.0 else None,
            )
            visualization_objects[handle_name] = handle
    except Exception as e:
        print(f"Warning: Could not visualize point {name}: {e}")


def points_viz(
    name: str,
    points: np.ndarray,
    radius: float = 0.01,
    color: float = 0xFF0000,
    opacity: float = 1.0,
) -> None:
    """
    Visualizes multiple points using viser's point cloud API for better performance
    """
    server = get_viewer()
    handle_name = f"points/{name}"

    # Convert hex color to RGB tuple
    r = (color >> 16) & 0xFF
    g = (color >> 8) & 0xFF
    b = color & 0xFF
    rgb_color = (r, g, b)

    try:
        # Ensure points is a numpy array with shape (N, 3)
        if len(points) == 0:
            # Hide or remove if no points
            if handle_name in visualization_objects:
                visualization_objects[handle_name].visible = False
            return

        points_array = np.asarray(points, dtype=np.float32)
        if points_array.ndim == 1:
            points_array = points_array.reshape(1, -1)
        if points_array.shape[1] != 3:
            raise ValueError(f"Points should have shape (N, 3), got {points_array.shape}")

        # Update or create point cloud
        if handle_name in visualization_objects:
            handle = visualization_objects[handle_name]
            handle.points = points_array
        else:
            handle = server.scene.add_point_cloud(
                handle_name,
                points=points_array,
                colors=rgb_color,
                point_size=radius,  # Convert radius to point_size (point_size is in screen pixels)
                point_shape="circle",
            )
            visualization_objects[handle_name] = handle

    except Exception as e:
        print(f"Warning: Could not visualize points {name}: {e}")


def robot_frame_viz(
    robot: placo.RobotWrapper, frame: str, opacity: float = 1.0, scale: float = 1.0
) -> None:
    """
    Draw a frame from the robot
    """
    node_name = f"{robot_names[robot]}_{frame}"
    frame_viz(node_name, robot.get_T_world_frame(frame), opacity, scale)


def line_viz(name: str, points: np.ndarray, color: float = 0xFF0000) -> None:
    """
    Visualizes a line using viser's add_line_segments API
    """
    server = get_viewer()
    handle_name = f"lines/{name}"

    # Convert hex color to RGB tuple
    r = (color >> 16) & 0xFF
    g = (color >> 8) & 0xFF
    b = color & 0xFF
    rgb_color = (r, g, b)

    try:
        if len(points) < 2:
            return

        # Prepare line segments: shape (N, 2, 3) where N is number of segments
        # Each segment has 2 points (start, end), each point is 3D
        segments = []
        for i in range(len(points) - 1):
            p1 = points[i]
            p2 = points[i + 1]
            # Skip zero-length segments
            if np.linalg.norm(p2 - p1) < 1e-6:
                continue
            segments.append([p1, p2])

        if len(segments) == 0:
            return

        # Convert to numpy array with shape (N, 2, 3)
        segments_array = np.array(segments, dtype=np.float32)

        # Update or create line segments
        if handle_name in visualization_objects:
            handle = visualization_objects[handle_name]
            handle.points = segments_array
        else:
            handle = server.scene.add_line_segments(
                handle_name,
                points=segments_array,
                colors=rgb_color,
                line_width=2.0,
            )
            visualization_objects[handle_name] = handle

    except Exception as e:
        print(f"Warning: Could not visualize line {name}: {e}")


def cross_viz(name: str, position: np.ndarray, size: float = 0.15, color: float = 0xFF0000) -> None:
    """
    Visualizes a cross
    """
    # Create two line segments forming a cross
    points1 = np.array([position - np.array([size, size, 0]), position + np.array([size, size, 0])])
    points2 = np.array(
        [position - np.array([size, -size, 0]), position + np.array([size, -size, 0])]
    )

    line_viz(f"{name}_1", points1, color)
    line_viz(f"{name}_2", points2, color)


def path_viz(name: str, points: np.ndarray, color: float = 0xFF0000) -> None:
    """
    Visualizes a path
    """
    key = f"path_{name}"

    # Delete old path segments
    server = get_viewer()
    for i in range(1000):  # Reasonable limit
        segment_name = f"lines/{key}_seg_{i}"
        if segment_name in visualization_objects:
            try:
                visualization_objects[segment_name].visible = False
                del visualization_objects[segment_name]
            except Exception:
                pass
        else:
            break

    # Create new path using line_viz which now uses add_line_segments (batched)
    line_viz(key, points, color)


def box_viz(
    name: str,
    position: np.ndarray,
    dimensions: np.ndarray,
    color: float = 0xFF0000,
    opacity: float = 1.0,
    wxyz: tuple[float, float, float, float] | np.ndarray = (1.0, 0.0, 0.0, 0.0),
) -> None:
    """
    Visualizes a box
    """
    server = get_viewer()
    handle_name = f"boxes/{name}"

    # Convert hex color to RGB tuple
    r = (color >> 16) & 0xFF
    g = (color >> 8) & 0xFF
    b = color & 0xFF
    rgb_color = (r, g, b)

    try:
        # Ensure dimensions is a numpy array
        if isinstance(dimensions, (list, tuple)):
            dimensions = np.array(dimensions)
        if dimensions.shape != (3,):
            raise ValueError(f"Dimensions must be a 3-element array, got {dimensions.shape}")

        if handle_name in visualization_objects:
            handle = visualization_objects[handle_name]
            handle.position = position
            handle.wxyz = wxyz
            handle.dimensions = tuple(dimensions)
        else:
            handle = server.scene.add_box(
                handle_name,
                color=rgb_color,
                dimensions=tuple(dimensions),
                position=position,
                wxyz=wxyz,
                opacity=opacity if opacity < 1.0 else None,
            )
            visualization_objects[handle_name] = handle
    except Exception as e:
        print(f"Warning: Could not visualize box {name}: {e}")


def cylinder_viz(
    name: str,
    position: np.ndarray,
    length: float,
    radius: float,
    color: float = 0xFF0000,
    opacity: float = 1.0,
) -> None:
    """
    Visualizes a cylinder using proper cylinder mesh
    """
    server = get_viewer()
    handle_name = f"cylinders/{name}"

    # Convert hex color to RGB tuple
    r = (color >> 16) & 0xFF
    g = (color >> 8) & 0xFF
    b = color & 0xFF
    rgb_color = (r, g, b)

    try:
        # Position cylinder with center at position, oriented along z-axis
        # Trimesh cylinder: base at z=0, top at z=height
        # We want center at position, so position at position - [0, 0, length/2]
        # Or position at position with cylinder extending from -length/2 to +length/2
        quaternion = (1.0, 0.0, 0.0, 0.0)  # No rotation (cylinder along z-axis)
        cylinder_center = position

        if HAS_TRIMESH:
            # Create proper cylinder mesh
            if handle_name in visualization_objects:
                visualization_objects[handle_name].remove()

            # Create cylinder with actual dimensions
            cylinder_mesh = trimesh.creation.cylinder(radius=radius, height=length, sections=16)
            # Translate so center is at z=0 (cylinder extends from -length/2 to +length/2)
            cylinder_mesh.apply_translation(np.array([0, 0, -length / 2]))

            handle = server.scene.add_mesh_simple(
                handle_name,
                cylinder_mesh.vertices,
                cylinder_mesh.faces,
                color=rgb_color,
                wxyz=quaternion,
                position=cylinder_center,
                opacity=opacity if opacity < 1.0 else None,
                cast_shadow=False,
                receive_shadow=False,
            )
            visualization_objects[handle_name] = handle
        else:
            # Fallback to box if trimesh is not available
            if handle_name in visualization_objects:
                handle = visualization_objects[handle_name]
                handle.position = cylinder_center
            else:
                handle = server.scene.add_box(
                    handle_name,
                    color=rgb_color,
                    dimensions=(radius * 2, radius * 2, length),
                    position=cylinder_center,
                    wxyz=quaternion,
                    opacity=opacity if opacity < 1.0 else None,
                )
                visualization_objects[handle_name] = handle
    except Exception as e:
        print(f"Warning: Could not visualize cylinder {name}: {e}")


def arrow_viz(
    name: str,
    point_from: np.ndarray,
    point_to: np.ndarray,
    color: float = 0xFF0000,
    radius: float = 0.003,
) -> None:
    """
    Visualizes an arrow using proper cylinder and cone meshes (similar to MuJoCo viser implementation)
    """
    server = get_viewer()

    # Arrow proportions: shaft is 80% of length, head is 20%
    # Ensure they add up exactly to total_length for seamless connection
    total_length = np.linalg.norm(point_to - point_from)
    if total_length < 1e-6:
        return

    shaft_length_ratio = 0.8
    head_length_ratio = 0.2
    # Calculate lengths ensuring they sum exactly to total_length
    shaft_length = total_length * shaft_length_ratio
    head_length = total_length * head_length_ratio
    # Ensure exact connection: shaft_length + head_length = total_length
    # Adjust head_length to account for any floating point errors
    head_length = total_length - shaft_length

    direction = (point_to - point_from) / total_length

    # Convert hex color to RGB tuple
    r = (color >> 16) & 0xFF
    g = (color >> 8) & 0xFF
    b = color & 0xFF
    rgb_color = (r, g, b)

    try:
        # Calculate rotation to align with direction (z-axis -> direction)
        z_axis = np.array([0, 0, 1])
        if np.abs(np.dot(direction, z_axis)) > 0.99:
            quaternion = (1.0, 0.0, 0.0, 0.0)
        else:
            axis = np.cross(z_axis, direction)
            axis = axis / (np.linalg.norm(axis) + 1e-6)
            angle = np.arccos(np.clip(np.dot(z_axis, direction), -1, 1))
            w = np.cos(angle / 2)
            s = np.sin(angle / 2)
            quaternion = (w, axis[0] * s, axis[1] * s, axis[2] * s)

        # Use proper meshes if trimesh is available, otherwise fall back to boxes
        if HAS_TRIMESH:
            # Create arrow shaft (cylinder)
            shaft_name = f"arrows/{name}/shaft"
            # Trimesh cylinder: base at z=0, top at z=height
            # After rotation, z-axis points along direction
            # Position at point_from so base (z=0) is at point_from
            # Top (z=shaft_length) will be at point_from + direction * shaft_length

            # Calculate where shaft ends (this is where cone base should be)
            shaft_end = point_from + direction * shaft_length

            # Create cylinder with actual dimensions
            shaft_mesh = trimesh.creation.cylinder(radius=radius, height=shaft_length, sections=16)

            if shaft_name in visualization_objects:
                # Remove old handle and create new one with updated mesh
                visualization_objects[shaft_name].remove()

            handle = server.scene.add_mesh_simple(
                shaft_name,
                shaft_mesh.vertices,
                shaft_mesh.faces,
                color=rgb_color,
                wxyz=quaternion,
                position=point_from,  # Base of cylinder at point_from
                cast_shadow=False,
                receive_shadow=False,
            )
            visualization_objects[shaft_name] = handle

            # Create arrow head (cone)
            head_name = f"arrows/{name}/head"
            # Trimesh cone: base at z=0, tip at z=height
            # After rotation, z-axis points along direction
            # Position cone so its base (z=0 in local coords) is exactly at shaft_end
            # This ensures seamless connection: shaft top (z=shaft_length) = cone base (z=0)
            head_radius = radius * 2.0

            # Create cone with actual dimensions
            # Trimesh cone: base center is at z=0, tip is at z=height
            head_mesh = trimesh.creation.cone(radius=head_radius, height=head_length, sections=8)

            # Position cone so its base (z=0) aligns exactly with shaft_end
            # Tip will be at shaft_end + direction * head_length = point_to
            # Verify: shaft_end + direction * head_length = point_from + direction * shaft_length + direction * head_length
            #         = point_from + direction * (shaft_length + head_length) = point_from + direction * total_length = point_to ✓
            head_position = shaft_end  # Base of cone at shaft end

            if head_name in visualization_objects:
                # Remove old handle and create new one with updated mesh
                visualization_objects[head_name].remove()

            handle = server.scene.add_mesh_simple(
                head_name,
                head_mesh.vertices,
                head_mesh.faces,
                color=rgb_color,
                wxyz=quaternion,
                position=head_position,  # Base of cone at shaft end
                cast_shadow=False,
                receive_shadow=False,
            )
            visualization_objects[head_name] = handle
        else:
            # Fallback to boxes if trimesh is not available
            shaft_position = point_from + direction * (shaft_length / 2)
            shaft_name = f"arrows/{name}/shaft"

            if shaft_name in visualization_objects:
                handle = visualization_objects[shaft_name]
                handle.position = shaft_position
                handle.wxyz = quaternion
            else:
                handle = server.scene.add_box(
                    shaft_name,
                    color=rgb_color,
                    dimensions=(radius * 2, radius * 2, shaft_length),
                    position=shaft_position,
                    wxyz=quaternion,
                )
                visualization_objects[shaft_name] = handle

            # Create arrow head (cone approximation using box)
            head_position = point_from + direction * (shaft_length + head_length / 2)
            head_name = f"arrows/{name}/head"

            if head_name in visualization_objects:
                handle = visualization_objects[head_name]
                handle.position = head_position
                handle.wxyz = quaternion
            else:
                # Use a larger box for the arrow head
                handle = server.scene.add_box(
                    head_name,
                    color=rgb_color,
                    dimensions=(radius * 4, radius * 4, head_length),
                    position=head_position,
                    wxyz=quaternion,
                )
                visualization_objects[head_name] = handle

    except Exception as e:
        print(f"Warning: Could not visualize arrow {name}: {e}")


previous_contacts: int = 0


def contacts_viz(solver: placo.DynamicsSolver, ratio: float = 0.1, radius: float = 0.005) -> None:
    """
    Visualizes contacts from a dynamics solver
    """
    global previous_contacts
    robot = solver.robot
    frames = robot.frame_names()
    k = 0

    for _ in range(solver.count_contacts()):
        contact = solver.get_contact(k)
        if not contact.active:
            continue

        if isinstance(contact, placo.PointContact):
            frame_name = frames[contact.position_task().frame_index]
            T_world_frame = robot.get_T_world_frame(frame_name)
            arrow_viz(
                f"contact_{k}",
                T_world_frame[:3, 3],
                T_world_frame[:3, 3] + contact.wrench * ratio,
                color=0x00FF00,
                radius=radius,
            )
        elif isinstance(contact, placo.Contact6D):
            frame_name = frames[contact.position_task().frame_index]
            T_world_frame = robot.get_T_world_frame(frame_name)
            wrench = T_world_frame[:3, :3] @ contact.wrench[:3]

            if np.linalg.norm(wrench) < 1e-6:
                origin = T_world_frame[:3, 3]
            else:
                origin = T_world_frame[:3, 3] + T_world_frame[:3, :3] @ contact.zmp()

            arrow_viz(
                f"contact_{k}",
                origin,
                origin + wrench * ratio,
                color=0x00FFAA,
                radius=radius,
            )
        elif isinstance(contact, placo.LineContact):
            frame_name = frames[contact.position_task().frame_index]
            T_world_frame = robot.get_T_world_frame(frame_name)
            wrench = T_world_frame[:3, :3] @ contact.wrench[:3]

            if np.linalg.norm(wrench) < 1e-6:
                origin = T_world_frame[:3, 3]
            else:
                origin = T_world_frame[:3, 3] + T_world_frame[:3, :3] @ contact.zmp()

            arrow_viz(
                f"contact_{k}",
                origin,
                origin + wrench * ratio,
                color=0xFF33AA,
                radius=radius,
            )
        elif isinstance(contact, placo.ExternalWrenchContact):
            frame_name = frames[contact.frame_index]
            T_world_frame = robot.get_T_world_frame(frame_name)
            arrow_viz(
                f"contact_{k}",
                T_world_frame[:3, 3],
                T_world_frame[:3, 3] + contact.w_ext[:3] * ratio,
                color=0xFF2222,
                radius=radius,
            )

        k += 1

    # Remove old contacts
    server = get_viewer()
    k_d = k
    while k_d < previous_contacts:
        shaft_name = f"arrows/contact_{k_d}/cylinder"
        head_name = f"arrows/contact_{k_d}/head"
        for contact_name in [shaft_name, head_name]:
            if contact_name in visualization_objects:
                try:
                    visualization_objects[contact_name].visible = False
                    del visualization_objects[contact_name]
                except Exception:
                    pass
        k_d += 1
    previous_contacts = k

I have writen a viser based visualization class, but it seems slower than meshcat . Still not figured why ?