tangger/lerobot
1
0
Fork 1
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

small fix unit tests

Browse Source
This commit is contained in:
Remi Cadene
2024-10-14 17:19:25 +02:00
parent e0394961cb
commit cde4945b29
3 changed files with 6 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
lerobot/common/robot_devices/control_utils.py
View File

@@ -316,7 +316,8 @@ def stop_recording(robot, listener, display_cameras):
robot.disconnect() robot.disconnect()
if not is_headless(): if not is_headless():
listener.stop() if listener is not None:
listener.stop()
if display_cameras: if display_cameras:
cv2.destroyAllWindows() cv2.destroyAllWindows()

3
lerobot/scripts/control_robot.py
View File

@@ -289,6 +289,9 @@ def record(
# Increment by one dataset["current_episode_index"] # Increment by one dataset["current_episode_index"]
save_current_episode(dataset) save_current_episode(dataset)
if events["stop_recording"]:
break
log_say("Stop recording", play_sounds, blocking=True) log_say("Stop recording", play_sounds, blocking=True)
stop_recording(robot, listener, display_cameras) stop_recording(robot, listener, display_cameras)

2
tests/test_control_robot.py
View File

@@ -446,5 +446,5 @@ def test_record_with_event_stop_recording(tmpdir, request, robot_type, mock, num
) )
assert not mock_events["exit_early"], "`exit_early` wasn't properly reset to False" assert not mock_events["exit_early"], "`exit_early` wasn't properly reset to False"
assert mock_add_frame.call_count == 1, "`add_frame` should have been called 2 times" assert mock_add_frame.call_count == 1, "`add_frame` should have been called 1 time"
assert len(dataset) == 1, "`dataset` should contain only 1 frame" assert len(dataset) == 1, "`dataset` should contain only 1 frame"