feat(vis): add raw action trajectory viewer
This commit is contained in:
Logic
@@ -0,0 +1,26 @@
|
||||
# Raw Action Trajectory Viewer Implementation Plan
|
||||
|
||||
> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.
|
||||
|
||||
**Goal:** 在可交互 MuJoCo 仿真窗口中，把 rollout 导出的 raw EE action 轨迹用红色轨迹标出来并启动仿真供人工查看。
|
||||
|
||||
**Architecture:** 读取已有 trajectory artifact 中的 raw_action / step 数据，生成左右臂末端轨迹点，并在 viewer 渲染循环中持续注入红色 marker。实现尽量独立为一个可复用的小脚本，避免影响训练/评估主路径。
|
||||
|
||||
**Tech Stack:** Python, NumPy, MuJoCo viewer, unittest/mock.
|
||||
|
||||
---
|
||||
|
||||
### Task 1: 抽取 raw_action 轨迹并生成可视化点集
|
||||
- [ ] 写失败测试，验证从 trajectory.npz 提取左右臂轨迹点
|
||||
- [ ] 实现最小 helper
|
||||
- [ ] 运行测试确认通过
|
||||
|
||||
### Task 2: 在 viewer 中渲染红色轨迹并支持交互查看
|
||||
- [ ] 写失败测试，验证 marker 配置/调用
|
||||
- [ ] 实现 viewer 可视化脚本
|
||||
- [ ] 运行测试确认通过
|
||||
|
||||
### Task 3: 启动真实仿真窗口供人工查看
|
||||
- [ ] 用现有 trajectory artifact 启动 viewer
|
||||
- [ ] 确认窗口可交互、红线出现
|
||||
- [ ] 向用户汇报启动方式与脚本路径
|
||||
36
roboimi/demos/view_raw_action_trajectory.py
Normal file
36
roboimi/demos/view_raw_action_trajectory.py
Normal file
@@ -0,0 +1,36 @@
|
||||
import argparse
|
||||
import numpy as np
|
||||
|
||||
from roboimi.utils.raw_action_trajectory_viewer import launch_raw_action_trajectory_viewer
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Launch an interactive MuJoCo viewer with raw-action trajectory overlay.")
|
||||
parser.add_argument("trajectory_path", help="Path to raw_action.npy or trajectory.npz")
|
||||
parser.add_argument("--task-name", default="sim_transfer")
|
||||
parser.add_argument("--line-radius", type=float, default=0.004)
|
||||
parser.add_argument("--max-markers", type=int, default=1500)
|
||||
parser.add_argument(
|
||||
"--box-pos",
|
||||
type=float,
|
||||
nargs=3,
|
||||
default=None,
|
||||
help="Optional box xyz to use when resetting the environment",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
box_pos = np.asarray(args.box_pos, dtype=np.float32) if args.box_pos is not None else None
|
||||
launch_raw_action_trajectory_viewer(
|
||||
args.trajectory_path,
|
||||
task_name=args.task_name,
|
||||
line_radius=args.line_radius,
|
||||
max_markers=args.max_markers,
|
||||
box_pos=box_pos,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
176
roboimi/utils/raw_action_trajectory_viewer.py
Normal file
176
roboimi/utils/raw_action_trajectory_viewer.py
Normal file
@@ -0,0 +1,176 @@
|
||||
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()
|
||||
119
tests/test_raw_action_trajectory_viewer.py
Normal file
119
tests/test_raw_action_trajectory_viewer.py
Normal file
@@ -0,0 +1,119 @@
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from unittest import mock
|
||||
|
||||
import numpy as np
|
||||
|
||||
from roboimi.utils import raw_action_trajectory_viewer as traj_view
|
||||
|
||||
|
||||
class RawActionTrajectoryViewerTest(unittest.TestCase):
|
||||
def test_set_transfer_box_pose_writes_joint_qpos(self):
|
||||
joint_qpos = np.zeros(7, dtype=np.float64)
|
||||
|
||||
class _FakeJoint:
|
||||
def __init__(self, qpos):
|
||||
self.qpos = qpos
|
||||
|
||||
class _FakeData:
|
||||
def joint(self, name):
|
||||
assert name == "red_box_joint"
|
||||
return _FakeJoint(joint_qpos)
|
||||
|
||||
traj_view.set_transfer_box_pose(_FakeData(), np.array([0.2, -0.1, 1.05], dtype=np.float64))
|
||||
|
||||
np.testing.assert_array_equal(
|
||||
joint_qpos,
|
||||
np.array([0.2, -0.1, 1.05, 1.0, 0.0, 0.0, 0.0], dtype=np.float64),
|
||||
)
|
||||
|
||||
def test_disable_cv2_highgui_temporarily_replaces_gui_calls(self):
|
||||
fake_cv2 = SimpleNamespace(
|
||||
namedWindow=lambda *args, **kwargs: "named",
|
||||
imshow=lambda *args, **kwargs: "imshow",
|
||||
waitKey=lambda *args, **kwargs: "wait",
|
||||
)
|
||||
|
||||
restore = traj_view.disable_cv2_highgui(fake_cv2)
|
||||
self.assertIsNone(fake_cv2.namedWindow("x"))
|
||||
self.assertIsNone(fake_cv2.imshow("x", None))
|
||||
self.assertEqual(fake_cv2.waitKey(1), 1)
|
||||
|
||||
restore()
|
||||
self.assertEqual(fake_cv2.namedWindow("x"), "named")
|
||||
self.assertEqual(fake_cv2.imshow("x", None), "imshow")
|
||||
self.assertEqual(fake_cv2.waitKey(1), "wait")
|
||||
|
||||
def test_load_raw_action_positions_from_npz(self):
|
||||
raw_action = np.array(
|
||||
[
|
||||
[1.0, 2.0, 3.0, 0, 0, 0, 1, 11.0, 12.0, 13.0, 0, 0, 0, 1, -1, -1],
|
||||
[4.0, 5.0, 6.0, 0, 0, 0, 1, 14.0, 15.0, 16.0, 0, 0, 0, 1, -1, -1],
|
||||
],
|
||||
dtype=np.float32,
|
||||
)
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
path = Path(tmpdir) / "trajectory.npz"
|
||||
np.savez(path, raw_action=raw_action)
|
||||
|
||||
positions = traj_view.load_raw_action_positions(path)
|
||||
|
||||
np.testing.assert_array_equal(
|
||||
positions["left"],
|
||||
np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=np.float32),
|
||||
)
|
||||
np.testing.assert_array_equal(
|
||||
positions["right"],
|
||||
np.array([[11.0, 12.0, 13.0], [14.0, 15.0, 16.0]], dtype=np.float32),
|
||||
)
|
||||
|
||||
def test_build_red_capsule_segments_downsamples_to_fit_scene_limit(self):
|
||||
left = np.stack([np.array([float(i), 0.0, 0.0], dtype=np.float32) for i in range(6)])
|
||||
right = np.stack([np.array([float(i), 1.0, 0.0], dtype=np.float32) for i in range(6)])
|
||||
|
||||
markers = traj_view.build_trajectory_capsule_markers(
|
||||
{"left": left, "right": right},
|
||||
max_markers=4,
|
||||
radius=0.01,
|
||||
)
|
||||
|
||||
self.assertLessEqual(len(markers), 4)
|
||||
self.assertTrue(all(marker["rgba"] == (1.0, 0.0, 0.0, 1.0) for marker in markers))
|
||||
self.assertTrue(all(marker["radius"] == 0.01 for marker in markers))
|
||||
|
||||
def test_apply_capsule_markers_populates_user_scene(self):
|
||||
fake_scene = SimpleNamespace(
|
||||
maxgeom=3,
|
||||
ngeom=99,
|
||||
geoms=[object(), object(), object()],
|
||||
)
|
||||
markers = [
|
||||
{
|
||||
"from": np.array([0.0, 0.0, 0.0], dtype=np.float64),
|
||||
"to": np.array([1.0, 0.0, 0.0], dtype=np.float64),
|
||||
"rgba": (1.0, 0.0, 0.0, 1.0),
|
||||
"radius": 0.01,
|
||||
},
|
||||
{
|
||||
"from": np.array([0.0, 1.0, 0.0], dtype=np.float64),
|
||||
"to": np.array([1.0, 1.0, 0.0], dtype=np.float64),
|
||||
"rgba": (1.0, 0.0, 0.0, 1.0),
|
||||
"radius": 0.01,
|
||||
},
|
||||
]
|
||||
|
||||
with mock.patch.object(traj_view.mujoco, "mjv_initGeom") as init_geom, mock.patch.object(
|
||||
traj_view.mujoco,
|
||||
"mjv_connector",
|
||||
) as connector:
|
||||
traj_view.apply_capsule_markers_to_scene(fake_scene, markers)
|
||||
|
||||
self.assertEqual(fake_scene.ngeom, 2)
|
||||
self.assertEqual(init_geom.call_count, 2)
|
||||
self.assertEqual(connector.call_count, 2)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user