From bf2fc099cc3373b6a8c1eff04d5a859e8765f9eb Mon Sep 17 00:00:00 2001 From: Remi Cadene Date: Fri, 26 Jul 2024 17:07:09 +0200 Subject: [PATCH] fix tests --- tests/test_control_robot.py | 23 ++++++++++++++--------- 1 file changed, 14 insertions(+), 9 deletions(-) diff --git a/tests/test_control_robot.py b/tests/test_control_robot.py index ae97fe1f..9ca303ed 100644 --- a/tests/test_control_robot.py +++ b/tests/test_control_robot.py @@ -3,13 +3,19 @@ from pathlib import Path from lerobot.common.policies.factory import make_policy from lerobot.common.robot_devices.robots.factory import make_robot from lerobot.common.utils.utils import init_hydra_config -from lerobot.scripts.control_robot import record_dataset, replay_episode, run_policy, teleoperate -from tests.utils import DEFAULT_CONFIG_PATH, DEVICE, require_koch +from lerobot.scripts.control_robot import record, replay, teleoperate +from tests.utils import DEFAULT_CONFIG_PATH, DEVICE, KOCH_ROBOT_CONFIG_PATH, require_koch + + +def make_robot_(): + robot_cfg = init_hydra_config(KOCH_ROBOT_CONFIG_PATH) + robot = make_robot(robot_cfg) + return robot @require_koch def test_teleoperate(request): - robot = make_robot("koch") + robot = make_robot_() teleoperate(robot, teleop_time_s=1) teleoperate(robot, fps=30, teleop_time_s=1) teleoperate(robot, fps=60, teleop_time_s=1) @@ -17,20 +23,19 @@ def test_teleoperate(request): @require_koch -def test_record_dataset_and_replay_episode_and_run_policy(tmpdir, request): - robot_name = "koch" +def test_record_and_replay_and_policy(tmpdir, request): env_name = "koch_real" policy_name = "act_koch_real" root = Path(tmpdir) repo_id = "lerobot/debug" - robot = make_robot(robot_name) - dataset = record_dataset( + robot = make_robot_() + dataset = record( robot, fps=30, root=root, repo_id=repo_id, warmup_time_s=1, episode_time_s=1, num_episodes=2 ) - replay_episode(robot, episode=0, fps=30, root=root, repo_id=repo_id) + replay(robot, episode=0, fps=30, root=root, repo_id=repo_id) cfg = init_hydra_config( DEFAULT_CONFIG_PATH, @@ -43,6 +48,6 @@ def test_record_dataset_and_replay_episode_and_run_policy(tmpdir, request): policy = make_policy(hydra_cfg=cfg, dataset_stats=dataset.stats) - run_policy(robot, policy, cfg, run_time_s=1) + record(robot, policy, cfg, run_time_s=1) del robot