274 lines
7.8 KiB
Python
274 lines
7.8 KiB
Python
import json
|
|
|
|
import pyarrow as pa
|
|
import pyarrow.parquet as parquet
|
|
import pytest
|
|
|
|
from lerobot_robot_ufactory.scripts import uf_lerobot_replay as replay_module
|
|
from lerobot_robot_ufactory.scripts.uf_lerobot_replay import (
|
|
JOINT_STATE_NAMES,
|
|
ReplayEpisode,
|
|
STATE_NAMES,
|
|
load_replay_episode,
|
|
replay_episode,
|
|
state_to_robot_action,
|
|
)
|
|
|
|
|
|
def _write_dataset(tmp_path, states, episode_indices=None, frame_indices=None, names=None):
|
|
dataset_root = tmp_path / "dataset"
|
|
(dataset_root / "data" / "chunk-000").mkdir(parents=True)
|
|
(dataset_root / "meta").mkdir()
|
|
|
|
episode_indices = episode_indices or [0] * len(states)
|
|
frame_indices = frame_indices or list(range(len(states)))
|
|
names = names or list(STATE_NAMES)
|
|
info = {
|
|
"fps": 30,
|
|
"features": {
|
|
"observation.state": {
|
|
"dtype": "float32",
|
|
"names": names,
|
|
"shape": [len(names)],
|
|
}
|
|
},
|
|
}
|
|
(dataset_root / "meta" / "info.json").write_text(
|
|
json.dumps(info), encoding="utf-8"
|
|
)
|
|
table = pa.table(
|
|
{
|
|
"observation.state": states,
|
|
"episode_index": episode_indices,
|
|
"frame_index": frame_indices,
|
|
}
|
|
)
|
|
parquet.write_table(table, dataset_root / "data" / "chunk-000" / "file-000.parquet")
|
|
return dataset_root
|
|
|
|
|
|
def test_load_replay_episode_sorts_and_maps_absolute_states(tmp_path):
|
|
state_a = [0.1, -0.2, 0.3, -0.4, 0.5, -0.6, 0.7, 1.00375]
|
|
state_b = [0.2, -0.3, 0.4, -0.5, 0.6, -0.7, 0.8, 0.4]
|
|
dataset_root = _write_dataset(
|
|
tmp_path,
|
|
states=[state_b, state_a, [9.0] * 8],
|
|
episode_indices=[0, 0, 1],
|
|
frame_indices=[1, 0, 0],
|
|
)
|
|
|
|
episode = load_replay_episode(dataset_root, episode_index=0)
|
|
|
|
assert episode.fps == 30
|
|
assert episode.frame_indices == (0, 1)
|
|
assert episode.states == (tuple(state_a), tuple(state_b))
|
|
assert state_to_robot_action(state_a) == dict(zip(STATE_NAMES, state_a, strict=True))
|
|
|
|
|
|
def test_replay_sends_absolute_values_without_delta_accumulation(monkeypatch):
|
|
class FakeRobot:
|
|
def __init__(self):
|
|
self.actions = []
|
|
|
|
def send_action(self, action):
|
|
self.actions.append(action.copy())
|
|
|
|
monkeypatch.setattr(replay_module, "precise_sleep", lambda _: None)
|
|
state_a = (0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8)
|
|
state_b = (0.2, 0.4, 0.6, 0.8, 1.0, 1.2, 1.4, 0.9)
|
|
episode = ReplayEpisode(
|
|
fps=30,
|
|
episode_index=0,
|
|
frame_indices=(0, 1),
|
|
states=(state_a, state_b),
|
|
)
|
|
robot = FakeRobot()
|
|
|
|
replay_episode(robot, episode, show_progress=False)
|
|
|
|
assert robot.actions == [
|
|
dict(zip(STATE_NAMES, state_a, strict=True)),
|
|
dict(zip(STATE_NAMES, state_b, strict=True)),
|
|
]
|
|
assert robot.actions[1]["J1.pos"] == state_b[0]
|
|
|
|
|
|
def test_ufactory_robot_sends_absolute_radian_joint_targets(monkeypatch):
|
|
from lerobot_robot_ufactory.robots.uf_robot import uf_robot as uf_robot_module
|
|
from lerobot_robot_ufactory.robots.uf_robot.uf_robot_config import UFRobotConfig
|
|
|
|
class FakeArm:
|
|
error_code = 0
|
|
mode = 6
|
|
|
|
def __init__(self):
|
|
self.calls = []
|
|
|
|
def set_mode(self, mode):
|
|
self.calls.append(("set_mode", mode))
|
|
self.mode = mode
|
|
return 0
|
|
|
|
def set_state(self, state):
|
|
self.calls.append(("set_state", state))
|
|
return 0
|
|
|
|
def set_servo_angle(self, **kwargs):
|
|
self.calls.append(("set_servo_angle", kwargs))
|
|
return 0
|
|
|
|
monkeypatch.setattr(uf_robot_module.time, "sleep", lambda _: None)
|
|
robot = uf_robot_module.UFRobot(
|
|
UFRobotConfig(
|
|
id="replay-test",
|
|
robot_dof=7,
|
|
control_space="joint",
|
|
gripper_type=0,
|
|
cameras={},
|
|
)
|
|
)
|
|
arm = FakeArm()
|
|
robot.real_arm = arm
|
|
robot._is_connected = True
|
|
target = state_to_robot_action((0.1, -0.2, 0.3, -0.4, 0.5, -0.6, 0.7, 0.8))
|
|
|
|
robot.send_action(target)
|
|
|
|
servo_calls = [kwargs for name, kwargs in arm.calls if name == "set_servo_angle"]
|
|
assert servo_calls == [
|
|
{
|
|
"angle": [0.1, -0.2, 0.3, -0.4, 0.5, -0.6, 0.7],
|
|
"speed": 0.2,
|
|
"is_radian": True,
|
|
"wait": True,
|
|
}
|
|
]
|
|
|
|
|
|
def test_reset_to_initial_enables_robot_before_motion(monkeypatch):
|
|
from lerobot_robot_ufactory.robots.uf_robot import uf_robot as uf_robot_module
|
|
from lerobot_robot_ufactory.robots.uf_robot.uf_robot_config import UFRobotConfig
|
|
|
|
class FakeArm:
|
|
mode = 0
|
|
state = 0
|
|
error_code = 0
|
|
warn_code = 0
|
|
|
|
def __init__(self):
|
|
self.calls = []
|
|
|
|
def motion_enable(self, enable=True):
|
|
self.calls.append(("motion_enable", enable))
|
|
return 0
|
|
|
|
def clean_error(self):
|
|
self.calls.append(("clean_error",))
|
|
return 0
|
|
|
|
def set_mode(self, mode):
|
|
self.calls.append(("set_mode", mode))
|
|
self.mode = mode
|
|
return 0
|
|
|
|
def set_state(self, state):
|
|
self.calls.append(("set_state", state))
|
|
self.state = state
|
|
return 0
|
|
|
|
def set_servo_angle(self, **kwargs):
|
|
self.calls.append(("set_servo_angle", kwargs))
|
|
return 0
|
|
|
|
monkeypatch.setattr(uf_robot_module.time, "sleep", lambda _: None)
|
|
robot = uf_robot_module.UFRobot(
|
|
UFRobotConfig(
|
|
id="replay-reset-test",
|
|
robot_dof=7,
|
|
control_space="joint",
|
|
gripper_type=0,
|
|
cameras={},
|
|
)
|
|
)
|
|
arm = FakeArm()
|
|
robot.real_arm = arm
|
|
robot._is_connected = True
|
|
robot._initial_point = [0.0] * 7
|
|
robot.configure = lambda: None
|
|
|
|
robot.reset_to_initial()
|
|
|
|
assert [name for name, *_ in arm.calls] == [
|
|
"motion_enable",
|
|
"clean_error",
|
|
"set_mode",
|
|
"set_state",
|
|
"set_servo_angle",
|
|
]
|
|
assert arm.calls[0] == ("motion_enable", True)
|
|
|
|
|
|
def test_ufactory_robot_replay_path_sends_absolute_servo_j_targets():
|
|
from lerobot_robot_ufactory.robots.uf_robot import uf_robot as uf_robot_module
|
|
from lerobot_robot_ufactory.robots.uf_robot.uf_robot_config import UFRobotConfig
|
|
|
|
class FakeArm:
|
|
error_code = 0
|
|
mode = 1
|
|
state = 0
|
|
warn_code = 0
|
|
|
|
def __init__(self):
|
|
self.calls = []
|
|
|
|
def set_servo_angle_j(self, angles, **kwargs):
|
|
self.calls.append(("set_servo_angle_j", angles, kwargs))
|
|
return 0
|
|
|
|
robot = uf_robot_module.UFRobot(
|
|
UFRobotConfig(
|
|
id="replay-servoj-test",
|
|
robot_dof=7,
|
|
control_space="joint",
|
|
joint_command_mode=1,
|
|
gripper_type=0,
|
|
cameras={},
|
|
)
|
|
)
|
|
arm = FakeArm()
|
|
robot.real_arm = arm
|
|
robot._is_connected = True
|
|
|
|
target = state_to_robot_action((0.1, -0.2, 0.3, -0.4, 0.5, -0.6, 0.7, 0.8))
|
|
robot.send_action(target)
|
|
|
|
assert arm.calls == [
|
|
(
|
|
"set_servo_angle_j",
|
|
[0.1, -0.2, 0.3, -0.4, 0.5, -0.6, 0.7],
|
|
{"speed": 0.2, "is_radian": True},
|
|
)
|
|
]
|
|
|
|
|
|
def test_load_replay_episode_rejects_wrong_state_schema(tmp_path):
|
|
states = [[0.0] * 8]
|
|
dataset_root = _write_dataset(tmp_path, states, names=list(JOINT_STATE_NAMES) + ["wrong"])
|
|
|
|
with pytest.raises(ValueError, match="absolute fields"):
|
|
load_replay_episode(dataset_root)
|
|
|
|
|
|
def test_load_replay_episode_rejects_non_contiguous_frames(tmp_path):
|
|
dataset_root = _write_dataset(tmp_path, [[0.0] * 8, [1.0] * 8], frame_indices=[0, 2])
|
|
|
|
with pytest.raises(ValueError, match="contiguous"):
|
|
load_replay_episode(dataset_root)
|
|
|
|
|
|
def test_state_to_robot_action_rejects_non_finite_values():
|
|
state = [0.0] * 7 + [float("nan")]
|
|
|
|
with pytest.raises(ValueError, match="non-finite"):
|
|
state_to_robot_action(state)
|