105 lines
3.3 KiB
Python
105 lines
3.3 KiB
Python
from __future__ import annotations
|
|
|
|
import math
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import h5py
|
|
import numpy as np
|
|
import pytest
|
|
|
|
TEST_DIR = Path(__file__).resolve().parent
|
|
if str(TEST_DIR) not in sys.path:
|
|
sys.path.insert(0, str(TEST_DIR))
|
|
|
|
import ik_method_comparison as comparison
|
|
|
|
|
|
def test_slerp_uses_shortest_arc_and_returns_unit_quaternion() -> None:
|
|
start = np.asarray([0.0, 0.0, 0.0, 1.0])
|
|
end = -np.asarray([0.0, 0.0, math.sin(0.1), math.cos(0.1)])
|
|
|
|
actual = comparison._slerp_quaternion(start, end, 0.5)
|
|
|
|
assert np.linalg.norm(actual) == pytest.approx(1.0)
|
|
assert actual == pytest.approx(
|
|
[0.0, 0.0, math.sin(0.05), math.cos(0.05)]
|
|
)
|
|
|
|
|
|
def test_resample_trajectory_keeps_endpoints_and_uses_requested_rate() -> None:
|
|
trajectory = comparison.EpisodeTrajectory(
|
|
source_path=Path("episode.hdf5"),
|
|
times_s=np.asarray([0.0, 0.5, 1.0]),
|
|
target_poses=np.asarray(
|
|
[
|
|
[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0],
|
|
[0.5, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0],
|
|
[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0],
|
|
]
|
|
),
|
|
initial_joints=np.zeros(7),
|
|
)
|
|
|
|
actual = comparison.resample_trajectory(trajectory, 4.0)
|
|
|
|
assert actual.times_s == pytest.approx([0.0, 0.25, 0.5, 0.75, 1.0])
|
|
assert actual.target_poses[0] == pytest.approx(trajectory.target_poses[0])
|
|
assert actual.target_poses[-1] == pytest.approx(trajectory.target_poses[-1])
|
|
assert actual.target_poses[:, 0] == pytest.approx(actual.times_s)
|
|
|
|
|
|
def _write_episode(path: Path) -> None:
|
|
poses = np.asarray(
|
|
[
|
|
[0.1, -0.2, 0.3, 0.0, 0.0, 0.0, 1.0],
|
|
[0.2, -0.2, 0.3, 0.0, 0.0, 0.0, 1.0],
|
|
[0.3, -0.2, 0.3, 0.0, 0.0, 0.0, 1.0],
|
|
[0.4, -0.2, 0.3, 0.0, 0.0, 0.0, 1.0],
|
|
],
|
|
dtype=np.float32,
|
|
)
|
|
with h5py.File(path, "w") as handle:
|
|
handle.attrs["arm"] = "right_rm75"
|
|
handle.attrs["pose_order"] = "x,y,z,qx,qy,qz,qw"
|
|
handle.create_dataset("debug/tcp/final_target_pose", data=poses)
|
|
handle.create_dataset(
|
|
"debug/timestamps/control_monotonic_ns",
|
|
data=np.asarray([0, 33_000_000, 66_000_000, 99_000_000]),
|
|
)
|
|
handle.create_dataset(
|
|
"debug/control/teleop_active", data=[0, 1, 1, 0]
|
|
)
|
|
handle.create_dataset(
|
|
"debug/control/action_valid", data=[1, 1, 1, 1]
|
|
)
|
|
handle.create_dataset(
|
|
"debug/control/command_sent", data=[0, 1, 1, 0]
|
|
)
|
|
qpos = np.zeros((4, 8), dtype=np.float32)
|
|
qpos[1, :7] = np.arange(7) * 0.1
|
|
handle.create_dataset("observations/qpos", data=qpos)
|
|
|
|
|
|
def test_load_episode_uses_longest_valid_run_and_first_valid_qpos(
|
|
tmp_path: Path,
|
|
) -> None:
|
|
path = tmp_path / "episode.hdf5"
|
|
_write_episode(path)
|
|
|
|
actual = comparison.load_episode(path)
|
|
|
|
assert actual.times_s == pytest.approx([0.0, 0.033])
|
|
assert actual.target_poses[:, 0] == pytest.approx([0.2, 0.3])
|
|
assert actual.initial_joints == pytest.approx(np.arange(7) * 0.1)
|
|
|
|
|
|
def test_load_episode_rejects_wrong_arm(tmp_path: Path) -> None:
|
|
path = tmp_path / "episode.hdf5"
|
|
_write_episode(path)
|
|
with h5py.File(path, "r+") as handle:
|
|
handle.attrs.modify("arm", "left_rm75")
|
|
|
|
with pytest.raises(ValueError, match="right_rm75"):
|
|
comparison.load_episode(path)
|