# Copyright [2021-2025] Thanh Nguyen
# Copyright [2022-2023] [CNRS, Toward SAS]
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Enhanced robot visualization utilities with better performance."""
from typing import List, Optional, Union
import numpy as np
import pinocchio as pin
from dataclasses import dataclass
[docs]
@dataclass
class VisualizationConfig:
"""Configuration for robot visualization."""
com_color: List[float] = None
axes_color: List[float] = None
force_color: List[float] = None
bbox_color: List[float] = None
scale_factor: float = 1.0
def __post_init__(self):
if self.com_color is None:
self.com_color = [1.0, 0.0, 0.0, 1.0] # Red
if self.axes_color is None:
self.axes_color = [1.0, 0.0, 0.0, 1.0] # Red
if self.force_color is None:
self.force_color = [1.0, 1.0, 0.0, 1.0] # Yellow
if self.bbox_color is None:
self.bbox_color = [0.5, 0.0, 0.5, 0.5] # Purple, transparent
[docs]
class RobotVisualizer:
"""Enhanced robot visualizer with better organization."""
def __init__(self, model: pin.Model, data: pin.Data, viz, config: Optional[VisualizationConfig] = None):
self.model = model
self.data = data
self.viz = viz
self.config = config or VisualizationConfig()
# Cache frequently used values
self.joint_names = model.names
self.nv = model.nv
[docs]
def update_kinematics(self, q: np.ndarray) -> None:
"""Update robot kinematics efficiently."""
pin.forwardKinematics(self.model, self.data, q)
pin.updateFramePlacements(self.model, self.data)
pin.centerOfMass(self.model, self.data, q, True)
pin.computeSubtreeMasses(self.model, self.data)
[docs]
def display_com(self, q: np.ndarray, frame_indices: List[int]) -> None:
"""Display center of mass positions with better naming."""
self.update_kinematics(q)
for i, frame_idx in enumerate(frame_indices[:-1]): # Exclude last frame
next_frame_idx = frame_indices[i + 1]
# Calculate link properties
link_length = np.linalg.norm(
self.data.oMf[next_frame_idx].translation -
self.data.oMf[frame_idx].translation
)
# Mass-based radius scaling
mass_ratio = self.data.mass[i] / self.data.mass[0] if self.data.mass[0] > 0 else 0.1
radius = link_length * mass_ratio * self.config.scale_factor
# COM placement
placement = self.data.oMf[frame_idx].copy()
parent_id = self.model.frames[frame_idx].parent - 1
if parent_id >= 0:
placement.translation = self.data.com[parent_id]
# Create and place sphere
sphere_name = f"world/com_sphere_{i}"
self.viz.viewer.gui.addSphere(sphere_name, radius, self.config.com_color)
self._apply_placement(sphere_name, placement)
[docs]
def display_axes(self, q: np.ndarray, axis_length: float = 0.15, axis_radius: float = 0.01) -> None:
"""Display coordinate axes for joints."""
self.update_kinematics(q)
for i, joint_name in enumerate(self.joint_names):
joint_id = self.model.getJointId(joint_name)
axis_name = f"world/axis_{joint_name}_{i}"
# Create axis visualization
self.viz.viewer.gui.addXYZaxis(
axis_name,
self.config.axes_color,
axis_radius * self.config.scale_factor,
axis_length * self.config.scale_factor
)
# Apply joint placement
self._apply_placement(axis_name, self.data.oMi[joint_id])
[docs]
def display_force(self, force: pin.Force, placement: pin.SE3,
scale: float = 1e-3, name: str = "force_vector") -> None:
"""Display force vector with better scaling."""
# Transform force to world frame
world_force = force.se3Action(placement)
force_magnitude = np.linalg.norm(world_force.linear)
if force_magnitude < 1e-10: # Skip negligible forces
return
# Calculate arrow properties
arrow_length = force_magnitude * scale * self.config.scale_factor
arrow_radius = arrow_length * 0.05
# Align arrow with force direction
force_direction = world_force.linear / force_magnitude
rotation = self._rotation_from_vectors([1, 0, 0], force_direction)
arrow_placement = placement.copy()
arrow_placement.rotation = rotation
# Create and place arrow
arrow_name = f"world/{name}_arrow"
self.viz.viewer.gui.addArrow(arrow_name, arrow_radius, arrow_length, self.config.force_color)
self._apply_placement(arrow_name, arrow_placement)
[docs]
def display_bounding_boxes(self, q: np.ndarray, com_min: np.ndarray,
com_max: np.ndarray, frame_indices: List[int]) -> None:
"""Display COM bounding boxes for optimization visualization."""
self.update_kinematics(q)
for i, frame_idx in enumerate(frame_indices):
# Get COM bounds for this segment
bounds_start = 3 * i
min_bounds = com_min[bounds_start:bounds_start + 3]
max_bounds = com_max[bounds_start:bounds_start + 3]
# Calculate box dimensions
box_size = (max_bounds - min_bounds) * self.config.scale_factor
# Box placement at COM location
placement = self.data.oMf[frame_idx].copy()
parent_id = self.model.frames[frame_idx].parent - 1
if parent_id >= 0:
placement.translation = self.data.com[parent_id]
# Create and place box
box_name = f"world/com_bbox_{i}"
self.viz.viewer.gui.addBox(
box_name, box_size[0], box_size[1], box_size[2], self.config.bbox_color
)
self._apply_placement(box_name, placement)
[docs]
def display_joints(self, q: np.ndarray) -> None:
"""Display joint frames efficiently."""
self.update_kinematics(q)
for i in range(self.nv):
joint_placement = self.data.oMi[i]
joint_name = f"world/joint_frame_{i}"
# Create joint frame visualization
self.viz.viewer.gui.addXYZaxis(
joint_name,
self.config.axes_color,
0.01 * self.config.scale_factor,
0.15 * self.config.scale_factor
)
# Apply configuration
self._apply_placement(joint_name, joint_placement)
def _apply_placement(self, name: str, placement: pin.SE3) -> None:
"""Apply SE3 placement to visualization object."""
self.viz.viewer.gui.applyConfiguration(name, pin.SE3ToXYZQUATtuple(placement))
self.viz.viewer.gui.refresh()
def _rotation_from_vectors(self, vec1: np.ndarray, vec2: np.ndarray) -> np.ndarray:
"""Calculate rotation matrix aligning vec1 to vec2."""
a = vec1 / np.linalg.norm(vec1)
b = vec2 / np.linalg.norm(vec2)
v = np.cross(a, b)
c = np.dot(a, b)
s = np.linalg.norm(v)
if s < 1e-10: # Vectors are parallel
return np.eye(3) if c > 0 else -np.eye(3)
# Rodrigues' rotation formula
vx = np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]])
return np.eye(3) + vx + vx @ vx * ((1 - c) / (s ** 2))
# Backward compatibility functions
[docs]
def place(viz, name: str, M: pin.SE3) -> None:
"""Legacy placement function."""
viz.viewer.gui.applyConfiguration(name, pin.SE3ToXYZQUATtuple(M))
[docs]
def display_COM(model: pin.Model, data: pin.Data, viz, q: np.ndarray, IDX: list) -> None:
"""Legacy COM display function."""
visualizer = RobotVisualizer(model, data, viz)
visualizer.display_com(q, IDX)
[docs]
def display_axes(model: pin.Model, data: pin.Data, viz, q: np.ndarray) -> None:
"""Legacy axes display function."""
visualizer = RobotVisualizer(model, data, viz)
visualizer.display_axes(q)
[docs]
def rotation_matrix_from_vectors(vec1: np.ndarray, vec2: np.ndarray) -> np.ndarray:
"""Legacy rotation function."""
visualizer = RobotVisualizer(None, None, None)
return visualizer._rotation_from_vectors(vec1, vec2)
[docs]
def display_force(viz, phi: pin.Force, M_se3: pin.SE3) -> None:
"""Legacy force display function."""
visualizer = RobotVisualizer(None, None, viz)
visualizer.display_force(phi, M_se3)
[docs]
def display_bounding_boxes(viz, model: pin.Model, data: pin.Data, q: np.ndarray,
COM_min: np.ndarray, COM_max: np.ndarray, IDX: list) -> None:
"""Legacy bounding box display function."""
visualizer = RobotVisualizer(model, data, viz)
visualizer.display_bounding_boxes(q, COM_min, COM_max, IDX)
[docs]
def display_joints(viz, model: pin.Model, data: pin.Data, q: np.ndarray) -> None:
"""Legacy joint display function."""
visualizer = RobotVisualizer(model, data, viz)
visualizer.display_joints(q)