roboimi/roboimi/utils/raw_action_trajectory_viewer.py

from __future__ import annotations

import math
import time
from pathlib import Path
from typing import Iterable

import cv2
import mujoco
import numpy as np

from roboimi.assets.robots.diana_med import BiDianaMed
from roboimi.envs.mujoco_base import MujocoEnv
from roboimi.envs.double_pos_ctrl_env import make_sim_env
from roboimi.utils.act_ex_utils import sample_transfer_pose


def _load_raw_action_array(path: str | Path) -> np.ndarray:
    path = Path(path)
    if path.suffix == ".npy":
        raw_action = np.load(path)
    elif path.suffix == ".npz":
        archive = np.load(path)
        if "raw_action" in archive:
            raw_action = archive["raw_action"]
        elif "raw_predicted_ee_action" in archive:
            raw_action = archive["raw_predicted_ee_action"]
        else:
            raise KeyError(f"{path} does not contain raw_action")
    else:
        raise ValueError(f"unsupported trajectory file: {path}")
    raw_action = np.asarray(raw_action, dtype=np.float32)
    if raw_action.ndim != 2 or raw_action.shape[1] < 10:
        raise ValueError(f"raw_action must have shape (T, 16)-like, got {raw_action.shape}")
    return raw_action


def disable_cv2_highgui(cv2_module=cv2):
    original = {
        "namedWindow": cv2_module.namedWindow,
        "imshow": cv2_module.imshow,
        "waitKey": cv2_module.waitKey,
    }

    cv2_module.namedWindow = lambda *args, **kwargs: None
    cv2_module.imshow = lambda *args, **kwargs: None
    cv2_module.waitKey = lambda *args, **kwargs: 1

    def restore():
        cv2_module.namedWindow = original["namedWindow"]
        cv2_module.imshow = original["imshow"]
        cv2_module.waitKey = original["waitKey"]

    return restore


def set_transfer_box_pose(mj_data, box_pos: np.ndarray) -> None:
    box_pos = np.asarray(box_pos, dtype=np.float64)
    if box_pos.shape != (3,):
        raise ValueError(f"box_pos must have shape (3,), got {box_pos.shape}")
    joint = mj_data.joint("red_box_joint")
    joint.qpos[0] = box_pos[0]
    joint.qpos[1] = box_pos[1]
    joint.qpos[2] = box_pos[2]
    joint.qpos[3] = 1.0
    joint.qpos[4] = 0.0
    joint.qpos[5] = 0.0
    joint.qpos[6] = 0.0


def load_raw_action_positions(path: str | Path) -> dict[str, np.ndarray]:
    raw_action = _load_raw_action_array(path)
    return {
        "left": raw_action[:, :3].astype(np.float32, copy=True),
        "right": raw_action[:, 7:10].astype(np.float32, copy=True),
    }


def _downsample_points(points: np.ndarray, stride: int) -> np.ndarray:
    sampled = points[::stride]
    if len(sampled) == 0:
        return points
    if not np.array_equal(sampled[-1], points[-1]):
        sampled = np.concatenate([sampled, points[-1:]], axis=0)
    return sampled


def build_trajectory_capsule_markers(
    positions: dict[str, np.ndarray],
    *,
    max_markers: int,
    radius: float = 0.003,
    rgba: tuple[float, float, float, float] = (1.0, 0.0, 0.0, 1.0),
) -> list[dict]:
    total_segments = sum(max(len(points) - 1, 0) for points in positions.values())
    if total_segments == 0:
        return []
    stride = max(1, math.ceil(total_segments / max_markers))
    markers = []
    for points in positions.values():
        sampled = _downsample_points(np.asarray(points, dtype=np.float64), stride)
        for idx in range(len(sampled) - 1):
            markers.append(
                {
                    "from": sampled[idx],
                    "to": sampled[idx + 1],
                    "rgba": rgba,
                    "radius": float(radius),
                }
            )
    return markers[:max_markers]


def apply_capsule_markers_to_scene(user_scn, markers: Iterable[dict]) -> None:
    user_scn.ngeom = 0
    for marker in markers:
        if user_scn.ngeom >= user_scn.maxgeom:
            break
        geom = user_scn.geoms[user_scn.ngeom]
        mujoco.mjv_initGeom(
            geom,
            mujoco.mjtGeom.mjGEOM_CAPSULE,
            np.zeros(3, dtype=np.float64),
            np.zeros(3, dtype=np.float64),
            np.eye(3, dtype=np.float64).reshape(-1),
            np.asarray(marker["rgba"], dtype=np.float32),
        )
        mujoco.mjv_connector(
            geom,
            mujoco.mjtGeom.mjGEOM_CAPSULE,
            float(marker["radius"]),
            np.asarray(marker["from"], dtype=np.float64),
            np.asarray(marker["to"], dtype=np.float64),
        )
        user_scn.ngeom += 1


def launch_raw_action_trajectory_viewer(
    trajectory_path: str | Path,
    *,
    task_name: str = "sim_transfer",
    line_radius: float = 0.004,
    max_markers: int = 1500,
    box_pos: np.ndarray | None = None,
    disable_camera_window: bool = True,
):
    positions = load_raw_action_positions(trajectory_path)
    if task_name != "sim_transfer":
        raise NotImplementedError(f"unsupported task_name: {task_name}")
    if box_pos is None:
        box_pos = sample_transfer_pose()

    robot = BiDianaMed()
    viewer_env = MujocoEnv(robot=robot, is_render=True, renderer="viewer", control_freq=30)
    viewer_env.reset()
    set_transfer_box_pose(viewer_env.mj_data, box_pos)
    mujoco.mj_forward(viewer_env.mj_model, viewer_env.mj_data)
    markers = build_trajectory_capsule_markers(
        positions,
        max_markers=max_markers,
        radius=line_radius,
    )

    if viewer_env.viewer is None or getattr(viewer_env.viewer, "user_scn", None) is None:
        raise RuntimeError("viewer does not expose user_scn; cannot render trajectory overlay")

    try:
        while viewer_env.viewer.is_running() and not viewer_env.exit_flag:
            with viewer_env.viewer.lock():
                apply_capsule_markers_to_scene(viewer_env.viewer.user_scn, markers)
            viewer_env.render()
            time.sleep(1 / 60.0)
    finally:
        viewer_env.exit_flag = True
        if getattr(viewer_env, "viewer", None) is not None:
            viewer_env.viewer.close()