feat(vis): add raw action trajectory viewer

tests/test_raw_action_trajectory_viewer.py

@@ -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()