gr00t-WholeBodyControl/gear_sonic/trl/utils/vis_scenepic.py

"""ScenePic skeleton visualizer. Generates interactive 3D HTML animations."""

import numpy as np
import torch
from typing import Union, List, Optional
from pathlib import Path

try:
    import scenepic as sp

    SCENEPIC_AVAILABLE = True
except ImportError:
    SCENEPIC_AVAILABLE = False


def _normalize_vec(v, eps=1e-8):
    return v / (torch.norm(v, dim=-1, keepdim=True) + eps)


def _quat_to_exp_map(q):
    """Quaternion (wxyz) to axis-angle."""
    w, xyz = q[..., 0:1], q[..., 1:4]
    sin_half = torch.norm(xyz, dim=-1, keepdim=True)
    angle = 2 * torch.atan2(sin_half, w)
    axis = xyz / (sin_half + 1e-8)
    return axis * angle


def _quat_between_two_vec(v1, v2, eps=1e-6):
    """Quaternion (wxyz) rotating v1 to v2."""
    v1, v2 = v1.reshape(-1, 3), v2.reshape(-1, 3)
    dot = (v1 * v2).sum(-1)
    cross = torch.cross(v1, v2, dim=-1)
    out = _normalize_vec(torch.cat([(1 + dot).unsqueeze(-1), cross], dim=-1))
    out[dot > 1 - eps] = torch.tensor([1.0, 0.0, 0.0, 0.0], device=v1.device)
    return out


def _make_floor_texture():
    from PIL import ImageColor

    c1 = np.tile(np.asarray(ImageColor.getcolor("#81C6EB", "RGB"), dtype=np.uint8), (5, 5, 1))
    c2 = np.tile(np.asarray(ImageColor.getcolor("#D4F1F7", "RGB"), dtype=np.uint8), (5, 5, 1))
    return np.tile(np.block([[[c1], [c2]], [[c2], [c1]]]), (15, 15, 1))


class SkeletonActor:
    """Skeleton with joints (spheres) and bones (cones)."""

    def __init__(
        self, scene, name, joint_parents, num_base_bodies=None, joint_radius=0.06, bone_radius=0.04
    ):
        self.joint_parents = joint_parents
        self.num_base = num_base_bodies or len(joint_parents)

        # Floor
        self.floor_img = scene.create_image(image_id="floor")
        self.floor_img.from_numpy(_make_floor_texture())
        self.floor_mesh = scene.create_mesh(texture_id="floor", layer_id="floor")
        self.floor_mesh.add_image(transform=sp.Transforms.Scale(20))

        # Joints and bones
        self.joint_meshes, self.bone_pairs = [], []
        for j, pa in enumerate(joint_parents):
            is_ext = j >= self.num_base
            jcolor = sp.Colors.Red if is_ext else (sp.Colors.Yellow if j == 0 else sp.Colors.Yellow)
            joint_mesh = scene.create_mesh(f"{name}_j{j}", layer_id=name)
            joint_mesh.add_sphere(color=jcolor, transform=sp.Transforms.scale(joint_radius))
            self.joint_meshes.append(joint_mesh)

            if pa >= 0:
                bcolor = sp.Colors.Orange if is_ext else sp.Colors.Green
                bone_mesh = scene.create_mesh(f"{name}_b{j}", layer_id=name)
                bone_mesh.add_cone(
                    color=bcolor,
                    transform=sp.Transforms.scale(np.array([1, bone_radius, bone_radius])),
                )
                self.bone_pairs.append((j, pa, bone_mesh))

    def add_to_frame(self, frame, pos):
        frame.add_mesh(self.floor_mesh)
        for j, p in enumerate(pos):
            frame.add_mesh(self.joint_meshes[j], transform=sp.Transforms.translate(p))

        if not self.bone_pairs:
            return
        vecs = np.stack([pos[j] - pos[pa] for j, pa, _ in self.bone_pairs])
        dists = np.linalg.norm(vecs, axis=-1)
        aa = _quat_to_exp_map(
            _quat_between_two_vec(
                torch.tensor([-1.0, 0.0, 0.0]).expand(len(vecs), 3),
                torch.tensor(vecs / (dists[:, None] + 1e-8)),
            )
        ).numpy()
        angles, axes = np.linalg.norm(aa, axis=-1, keepdims=True), aa / (
            np.linalg.norm(aa, axis=-1, keepdims=True) + 1e-8
        )

        for (j, pa, mesh), ang, ax, d in zip(self.bone_pairs, angles, axes, dists):
            t = sp.Transforms.translate((pos[pa] + pos[j]) * 0.5)
            t = (
                t
                @ sp.Transforms.RotationMatrixFromAxisAngle(ax, ang)
                @ sp.Transforms.Scale(np.array([d, 1, 1]))
            )
            frame.add_mesh(mesh, transform=t)


class ScenepicVisualizer:
    """Interactive 3D skeleton visualizer using ScenePic."""

    def __init__(self, joint_parents, num_base_bodies=None):
        if not SCENEPIC_AVAILABLE:
            raise ImportError("pip install scenepic")
        if isinstance(joint_parents, torch.Tensor):
            joint_parents = joint_parents.cpu().numpy().tolist()
        self.joint_parents = joint_parents
        self.num_base = num_base_bodies or len(joint_parents)

    def visualize(self, joint_positions, output_path, fps=30, title="Skeleton"):
        """Render skeleton animation to HTML. joint_positions: [T, J, 3]"""
        if isinstance(joint_positions, torch.Tensor):
            joint_positions = joint_positions.cpu().numpy()
        if joint_positions.ndim == 4:
            joint_positions = joint_positions[0]

        T = joint_positions.shape[0]
        Path(output_path).parent.mkdir(parents=True, exist_ok=True)

        scene = sp.Scene()
        scene.framerate = fps
        canvas = scene.create_canvas_3d(width=600, height=600)
        canvas.set_layer_settings({title: {"filled": True}})

        skeleton = SkeletonActor(scene, title, self.joint_parents, self.num_base)
        camera = sp.Camera(
            center=(4, 0, 1.5), look_at=(0, 0, 0.8), up_dir=(0, 0, 1), fov_y_degrees=45.0
        )

        for t in range(T):
            frame = canvas.create_frame()
            frame.camera = camera
            skeleton.add_to_frame(frame, joint_positions[t])

        scene.save_as_html(str(output_path))
        print(f"Saved: {output_path}")
        return output_path


if __name__ == "__main__":
    from pathlib import Path
    from omegaconf import OmegaConf
    from groot.rl.utils.motion_lib.torch_humanoid_batch import Humanoid_Batch
    from groot.rl.trl.utils.mujoco_fk_utils import MuJoCoFKHelper, load_qpos_from_csv

    groot_root = Path(__file__).parent.parent.parent.parent
    motion_yaml = (
        groot_root
        / "rl"
        / "config"
        / "manager_env"
        / "commands"
        / "terms"
        / "motion_g1_extended_toe.yaml"
    )
    cfg = OmegaConf.load(motion_yaml).motion.motion_lib_cfg
    fk_helper = MuJoCoFKHelper(Humanoid_Batch(cfg, device=torch.device("cpu")))
    print(
        f"Loaded: {fk_helper.num_bodies} + {fk_helper.num_bodies_augment - fk_helper.num_bodies} extended bodies"
    )

    csv_path = groot_root / ".." / "data" / "example_csv_g1_navigation.csv"
    if csv_path.exists():
        qpos = load_qpos_from_csv(str(csv_path))
    else:
        T, t = 60, torch.linspace(0, 4 * 3.14159, 60)
        dof = torch.zeros(T, fk_helper.num_dof)
        dof[:, 0], dof[:, 3] = 0.3 * torch.sin(t), 0.5 * torch.sin(t)
        dof[:, 6], dof[:, 9] = 0.3 * torch.sin(t + 3.14), 0.5 * torch.sin(t + 3.14)
        qpos = torch.cat(
            [
                torch.zeros(T, 3) + torch.tensor([0, 0, 1.0]),
                torch.tensor([[1.0, 0.0, 0.0, 0.0]]).expand(T, 4),
                dof,
            ],
            dim=-1,
        )

    global_pos, _ = fk_helper.qpos_to_global_transforms(
        qpos.unsqueeze(0), False, include_extended=True
    )

    viz = ScenepicVisualizer(fk_helper._parents, fk_helper.num_bodies)
    viz.visualize(
        global_pos[0],
        Path(__file__).parent.parent.parent.parent / "scenepic_demo.html",
        title="G1 Robot",
    )
    print("Yellow = base joints, Red = extended (head/toes)")