diff --git a/docs/source/envhub_leisaac.mdx b/docs/source/envhub_leisaac.mdx index 5c951b564..ff848d415 100644 --- a/docs/source/envhub_leisaac.mdx +++ b/docs/source/envhub_leisaac.mdx @@ -139,7 +139,7 @@ from lerobot.teleoperators import ( # noqa: F401 make_teleoperator_from_config, so101_leader, ) -from lerobot.utils.robot_utils import busy_wait +from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.utils import init_logging from lerobot.envs.factory import make_env @@ -196,7 +196,7 @@ def teleop_loop(teleop: Teleoperator, env: gym.Env, fps: int): obs, info = env.reset() dt_s = time.perf_counter() - loop_start - busy_wait(1 / fps - dt_s) + precise_sleep(1 / fps - dt_s) loop_s = time.perf_counter() - loop_start print(f"\ntime: {loop_s * 1e3:.2f}ms ({1 / loop_s:.0f} Hz)") diff --git a/docs/source/il_robots.mdx b/docs/source/il_robots.mdx index 000df2a19..a7b4fe570 100644 --- a/docs/source/il_robots.mdx +++ b/docs/source/il_robots.mdx @@ -393,7 +393,7 @@ import time from lerobot.datasets.lerobot_dataset import LeRobotDataset from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig from lerobot.robots.so100_follower.so100_follower import SO100Follower -from lerobot.utils.robot_utils import busy_wait +from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.utils import log_say episode_idx = 0 @@ -415,7 +415,7 @@ for idx in range(dataset.num_frames): } robot.send_action(action) - busy_wait(1.0 / dataset.fps - (time.perf_counter() - t0)) + precise_sleep(1.0 / dataset.fps - (time.perf_counter() - t0)) robot.disconnect() ``` diff --git a/examples/backward_compatibility/replay.py b/examples/backward_compatibility/replay.py index 6bca0570f..ed52a24c9 100644 --- a/examples/backward_compatibility/replay.py +++ b/examples/backward_compatibility/replay.py @@ -45,7 +45,7 @@ from lerobot.robots import ( # noqa: F401 so101_follower, ) from lerobot.utils.constants import ACTION -from lerobot.utils.robot_utils import busy_wait +from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.utils import ( init_logging, log_say, @@ -97,7 +97,7 @@ def replay(cfg: ReplayConfig): robot.send_action(action) dt_s = time.perf_counter() - start_episode_t - busy_wait(1 / dataset.fps - dt_s) + precise_sleep(1 / dataset.fps - dt_s) robot.disconnect() diff --git a/examples/lekiwi/replay.py b/examples/lekiwi/replay.py index 38cac20a0..872dacf27 100644 --- a/examples/lekiwi/replay.py +++ b/examples/lekiwi/replay.py @@ -20,7 +20,7 @@ from lerobot.datasets.lerobot_dataset import LeRobotDataset from lerobot.robots.lekiwi.config_lekiwi import LeKiwiClientConfig from lerobot.robots.lekiwi.lekiwi_client import LeKiwiClient from lerobot.utils.constants import ACTION -from lerobot.utils.robot_utils import busy_wait +from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.utils import log_say EPISODE_IDX = 0 @@ -58,7 +58,7 @@ def main(): # Send action to robot _ = robot.send_action(action) - busy_wait(max(1.0 / dataset.fps - (time.perf_counter() - t0), 0.0)) + precise_sleep(max(1.0 / dataset.fps - (time.perf_counter() - t0), 0.0)) robot.disconnect() diff --git a/examples/lekiwi/teleoperate.py b/examples/lekiwi/teleoperate.py index 870a99c65..c4d20ebbe 100644 --- a/examples/lekiwi/teleoperate.py +++ b/examples/lekiwi/teleoperate.py @@ -19,7 +19,7 @@ import time from lerobot.robots.lekiwi import LeKiwiClient, LeKiwiClientConfig from lerobot.teleoperators.keyboard.teleop_keyboard import KeyboardTeleop, KeyboardTeleopConfig from lerobot.teleoperators.so100_leader import SO100Leader, SO100LeaderConfig -from lerobot.utils.robot_utils import busy_wait +from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.visualization_utils import init_rerun, log_rerun_data FPS = 30 @@ -71,7 +71,7 @@ def main(): # Visualize log_rerun_data(observation=observation, action=action) - busy_wait(max(1.0 / FPS - (time.perf_counter() - t0), 0.0)) + precise_sleep(max(1.0 / FPS - (time.perf_counter() - t0), 0.0)) if __name__ == "__main__": diff --git a/examples/phone_to_so100/replay.py b/examples/phone_to_so100/replay.py index 21b37de8e..a7b18a53c 100644 --- a/examples/phone_to_so100/replay.py +++ b/examples/phone_to_so100/replay.py @@ -29,7 +29,7 @@ from lerobot.robots.so100_follower.robot_kinematic_processor import ( ) from lerobot.robots.so100_follower.so100_follower import SO100Follower from lerobot.utils.constants import ACTION -from lerobot.utils.robot_utils import busy_wait +from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.utils import log_say EPISODE_IDX = 0 @@ -96,7 +96,7 @@ def main(): # Send action to robot _ = robot.send_action(joint_action) - busy_wait(1.0 / dataset.fps - (time.perf_counter() - t0)) + precise_sleep(1.0 / dataset.fps - (time.perf_counter() - t0)) # Clean up robot.disconnect() diff --git a/examples/phone_to_so100/teleoperate.py b/examples/phone_to_so100/teleoperate.py index b9f8d3ad9..2ac8b3cce 100644 --- a/examples/phone_to_so100/teleoperate.py +++ b/examples/phone_to_so100/teleoperate.py @@ -32,7 +32,7 @@ from lerobot.robots.so100_follower.so100_follower import SO100Follower from lerobot.teleoperators.phone.config_phone import PhoneConfig, PhoneOS from lerobot.teleoperators.phone.phone_processor import MapPhoneActionToRobotAction from lerobot.teleoperators.phone.teleop_phone import Phone -from lerobot.utils.robot_utils import busy_wait +from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.visualization_utils import init_rerun, log_rerun_data FPS = 30 @@ -114,7 +114,7 @@ def main(): # Visualize log_rerun_data(observation=phone_obs, action=joint_action) - busy_wait(max(1.0 / FPS - (time.perf_counter() - t0), 0.0)) + precise_sleep(max(1.0 / FPS - (time.perf_counter() - t0), 0.0)) if __name__ == "__main__": diff --git a/examples/so100_to_so100_EE/replay.py b/examples/so100_to_so100_EE/replay.py index 0fc36f776..9951b139d 100644 --- a/examples/so100_to_so100_EE/replay.py +++ b/examples/so100_to_so100_EE/replay.py @@ -30,7 +30,7 @@ from lerobot.robots.so100_follower.robot_kinematic_processor import ( ) from lerobot.robots.so100_follower.so100_follower import SO100Follower from lerobot.utils.constants import ACTION -from lerobot.utils.robot_utils import busy_wait +from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.utils import log_say EPISODE_IDX = 0 @@ -97,7 +97,7 @@ def main(): # Send action to robot _ = robot.send_action(joint_action) - busy_wait(1.0 / dataset.fps - (time.perf_counter() - t0)) + precise_sleep(1.0 / dataset.fps - (time.perf_counter() - t0)) # Clean up robot.disconnect() diff --git a/examples/so100_to_so100_EE/teleoperate.py b/examples/so100_to_so100_EE/teleoperate.py index 3365f5a56..21299103b 100644 --- a/examples/so100_to_so100_EE/teleoperate.py +++ b/examples/so100_to_so100_EE/teleoperate.py @@ -32,7 +32,7 @@ from lerobot.robots.so100_follower.robot_kinematic_processor import ( from lerobot.robots.so100_follower.so100_follower import SO100Follower from lerobot.teleoperators.so100_leader.config_so100_leader import SO100LeaderConfig from lerobot.teleoperators.so100_leader.so100_leader import SO100Leader -from lerobot.utils.robot_utils import busy_wait +from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.visualization_utils import init_rerun, log_rerun_data FPS = 30 @@ -120,7 +120,7 @@ def main(): # Visualize log_rerun_data(observation=leader_ee_act, action=follower_joints_act) - busy_wait(max(1.0 / FPS - (time.perf_counter() - t0), 0.0)) + precise_sleep(max(1.0 / FPS - (time.perf_counter() - t0), 0.0)) if __name__ == "__main__": diff --git a/src/lerobot/rl/actor.py b/src/lerobot/rl/actor.py index 54d0fba69..13fd66507 100644 --- a/src/lerobot/rl/actor.py +++ b/src/lerobot/rl/actor.py @@ -78,7 +78,7 @@ from lerobot.transport.utils import ( transitions_to_bytes, ) from lerobot.utils.random_utils import set_seed -from lerobot.utils.robot_utils import busy_wait +from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.transition import ( Transition, move_state_dict_to_device, @@ -398,7 +398,7 @@ def act_with_policy( if cfg.env.fps is not None: dt_time = time.perf_counter() - start_time - busy_wait(1 / cfg.env.fps - dt_time) + precise_sleep(1 / cfg.env.fps - dt_time) # Communication Functions - Group all gRPC/messaging functions diff --git a/src/lerobot/rl/gym_manipulator.py b/src/lerobot/rl/gym_manipulator.py index f9c9d0d7a..ad1fdf55f 100644 --- a/src/lerobot/rl/gym_manipulator.py +++ b/src/lerobot/rl/gym_manipulator.py @@ -74,7 +74,7 @@ from lerobot.teleoperators import ( from lerobot.teleoperators.teleoperator import Teleoperator from lerobot.teleoperators.utils import TeleopEvents from lerobot.utils.constants import ACTION, DONE, OBS_IMAGES, OBS_STATE, REWARD -from lerobot.utils.robot_utils import busy_wait +from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.utils import log_say logging.basicConfig(level=logging.INFO) @@ -114,7 +114,7 @@ def reset_follower_position(robot_arm: Robot, target_position: np.ndarray) -> No for pose in trajectory: action_dict = dict(zip(current_position_dict, pose, strict=False)) robot_arm.bus.sync_write("Goal_Position", action_dict) - busy_wait(0.015) + precise_sleep(0.015) class RobotEnv(gym.Env): @@ -238,7 +238,7 @@ class RobotEnv(gym.Env): reset_follower_position(self.robot, np.array(self.reset_pose)) log_say("Reset the environment done.", play_sounds=True) - busy_wait(self.reset_time_s - (time.perf_counter() - start_time)) + precise_sleep(self.reset_time_s - (time.perf_counter() - start_time)) super().reset(seed=seed, options=options) @@ -713,7 +713,7 @@ def control_loop( transition = env_processor(transition) # Maintain fps timing - busy_wait(dt - (time.perf_counter() - step_start_time)) + precise_sleep(dt - (time.perf_counter() - step_start_time)) if dataset is not None and cfg.dataset.push_to_hub: logging.info("Pushing dataset to hub") @@ -745,7 +745,7 @@ def replay_trajectory( ) transition = action_processor(transition) env.step(transition[TransitionKey.ACTION]) - busy_wait(1 / cfg.env.fps - (time.perf_counter() - start_time)) + precise_sleep(1 / cfg.env.fps - (time.perf_counter() - start_time)) @parser.wrap() diff --git a/src/lerobot/scripts/lerobot_find_joint_limits.py b/src/lerobot/scripts/lerobot_find_joint_limits.py index 07d57a760..4ea83c976 100644 --- a/src/lerobot/scripts/lerobot_find_joint_limits.py +++ b/src/lerobot/scripts/lerobot_find_joint_limits.py @@ -50,7 +50,7 @@ from lerobot.teleoperators import ( # noqa: F401 make_teleoperator_from_config, so100_leader, ) -from lerobot.utils.robot_utils import busy_wait +from lerobot.utils.robot_utils import precise_sleep @dataclass @@ -114,7 +114,7 @@ def find_joint_and_ee_bounds(cfg: FindJointLimitsConfig): print(f"Min joint pos position {np.round(min_pos, 4).tolist()}") break - busy_wait(0.01) + precise_sleep(0.01) def main(): diff --git a/src/lerobot/scripts/lerobot_record.py b/src/lerobot/scripts/lerobot_record.py index 6df92d893..600a73022 100644 --- a/src/lerobot/scripts/lerobot_record.py +++ b/src/lerobot/scripts/lerobot_record.py @@ -119,7 +119,7 @@ from lerobot.utils.control_utils import ( sanity_check_dataset_robot_compatibility, ) from lerobot.utils.import_utils import register_third_party_devices -from lerobot.utils.robot_utils import busy_wait +from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.utils import ( get_safe_torch_device, init_logging, @@ -364,7 +364,7 @@ def record_loop( log_rerun_data(observation=obs_processed, action=action_values) dt_s = time.perf_counter() - start_loop_t - busy_wait(1 / fps - dt_s) + precise_sleep(1 / fps - dt_s) timestamp = time.perf_counter() - start_episode_t diff --git a/src/lerobot/scripts/lerobot_replay.py b/src/lerobot/scripts/lerobot_replay.py index ffd7b2b22..d25fcc3bd 100644 --- a/src/lerobot/scripts/lerobot_replay.py +++ b/src/lerobot/scripts/lerobot_replay.py @@ -62,7 +62,7 @@ from lerobot.robots import ( # noqa: F401 ) from lerobot.utils.constants import ACTION from lerobot.utils.import_utils import register_third_party_devices -from lerobot.utils.robot_utils import busy_wait +from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.utils import ( init_logging, log_say, @@ -121,7 +121,7 @@ def replay(cfg: ReplayConfig): _ = robot.send_action(processed_action) dt_s = time.perf_counter() - start_episode_t - busy_wait(1 / dataset.fps - dt_s) + precise_sleep(1 / dataset.fps - dt_s) robot.disconnect() diff --git a/src/lerobot/scripts/lerobot_teleoperate.py b/src/lerobot/scripts/lerobot_teleoperate.py index 0a418f3bc..ad35242b9 100644 --- a/src/lerobot/scripts/lerobot_teleoperate.py +++ b/src/lerobot/scripts/lerobot_teleoperate.py @@ -89,7 +89,7 @@ from lerobot.teleoperators import ( # noqa: F401 so101_leader, ) from lerobot.utils.import_utils import register_third_party_devices -from lerobot.utils.robot_utils import busy_wait +from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.utils import init_logging, move_cursor_up from lerobot.utils.visualization_utils import init_rerun, log_rerun_data @@ -173,7 +173,7 @@ def teleop_loop( move_cursor_up(len(robot_action_to_send) + 5) dt_s = time.perf_counter() - loop_start - busy_wait(1 / fps - dt_s) + precise_sleep(1 / fps - dt_s) loop_s = time.perf_counter() - loop_start print(f"\ntime: {loop_s * 1e3:.2f}ms ({1 / loop_s:.0f} Hz)") diff --git a/src/lerobot/utils/robot_utils.py b/src/lerobot/utils/robot_utils.py index 42abcdda4..28c8e7c49 100644 --- a/src/lerobot/utils/robot_utils.py +++ b/src/lerobot/utils/robot_utils.py @@ -16,14 +16,40 @@ import platform import time -def busy_wait(seconds): - if platform.system() == "Darwin" or platform.system() == "Windows": - # On Mac and Windows, `time.sleep` is not accurate and we need to use this while loop trick, - # but it consumes CPU cycles. +def precise_sleep(seconds: float, spin_threshold: float = 0.010, sleep_margin: float = 0.003): + """ + Wait for `seconds` with better precision than time.sleep alone at the expense of more CPU usage. + + Parameters: + - seconds: duration to wait + - spin_threshold: if remaining <= spin_threshold -> spin; otherwise sleep (seconds). Default 10ms + - sleep_margin: when sleeping leave this much time before deadline to avoid oversleep. Default 3ms + + Note: + The default parameters are chosen to prioritize timing accuracy over CPU usage for the common 30 FPS use case. + """ + if seconds <= 0: + return + + system = platform.system() + # On macOS and Windows the scheduler / sleep granularity can make + # short sleeps inaccurate. Instead of burning CPU for the whole + # duration, sleep for most of the time and spin for the final few + # milliseconds to achieve good accuracy with much lower CPU usage. + if system in ("Darwin", "Windows"): end_time = time.perf_counter() + seconds - while time.perf_counter() < end_time: - pass + while True: + remaining = end_time - time.perf_counter() + if remaining <= 0: + break + # If there's more than a couple milliseconds left, sleep most + # of the remaining time and leave a small margin for the final spin. + if remaining > spin_threshold: + # Sleep but avoid sleeping past the end by leaving a small margin. + time.sleep(max(remaining - sleep_margin, 0)) + else: + # Final short spin to hit precise timing without long sleeps. + pass else: - # On Linux time.sleep is accurate - if seconds > 0: - time.sleep(seconds) + # On Linux time.sleep is accurate enough for most uses + time.sleep(seconds)