diff --git a/docs/source/il_robots.mdx b/docs/source/il_robots.mdx index 19b62167e..91df14028 100644 --- a/docs/source/il_robots.mdx +++ b/docs/source/il_robots.mdx @@ -200,7 +200,7 @@ from lerobot.teleoperators.so100_leader.config_so100_leader import SO100LeaderCo from lerobot.teleoperators.so100_leader.so100_leader import SO100Leader from lerobot.utils.control_utils import init_keyboard_listener from lerobot.utils.utils import log_say -from lerobot.utils.visualization_utils import _init_rerun +from lerobot.utils.visualization_utils import init_rerun from lerobot.record import record_loop NUM_EPISODES = 5 @@ -237,7 +237,7 @@ dataset = LeRobotDataset.create( # Initialize the keyboard listener and rerun visualization _, events = init_keyboard_listener() -_init_rerun(session_name="recording") +init_rerun(session_name="recording") # Connect the robot and teleoperator robot.connect() @@ -517,7 +517,7 @@ from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerCon from lerobot.robots.so100_follower.so100_follower import SO100Follower from lerobot.utils.control_utils import init_keyboard_listener from lerobot.utils.utils import log_say -from lerobot.utils.visualization_utils import _init_rerun +from lerobot.utils.visualization_utils import init_rerun from lerobot.record import record_loop from lerobot.policies.factory import make_processor @@ -557,7 +557,7 @@ dataset = LeRobotDataset.create( # Initialize the keyboard listener and rerun visualization _, events = init_keyboard_listener() -_init_rerun(session_name="recording") +init_rerun(session_name="recording") # Connect the robot robot.connect() diff --git a/examples/lekiwi/evaluate.py b/examples/lekiwi/evaluate.py index 8993a5e14..32a5e0a2b 100644 --- a/examples/lekiwi/evaluate.py +++ b/examples/lekiwi/evaluate.py @@ -23,7 +23,7 @@ from lerobot.robots.lekiwi import LeKiwiClient, LeKiwiClientConfig from lerobot.scripts.lerobot_record import record_loop from lerobot.utils.control_utils import init_keyboard_listener from lerobot.utils.utils import log_say -from lerobot.utils.visualization_utils import _init_rerun +from lerobot.utils.visualization_utils import init_rerun NUM_EPISODES = 2 FPS = 30 @@ -73,7 +73,7 @@ teleop_action_processor, robot_action_processor, robot_observation_processor = m # Initialize the keyboard listener and rerun visualization listener, events = init_keyboard_listener() -_init_rerun(session_name="lekiwi_evaluate") +init_rerun(session_name="lekiwi_evaluate") if not robot.is_connected: raise ValueError("Robot is not connected!") diff --git a/examples/lekiwi/record.py b/examples/lekiwi/record.py index f59093b26..30f34e718 100644 --- a/examples/lekiwi/record.py +++ b/examples/lekiwi/record.py @@ -24,7 +24,7 @@ from lerobot.teleoperators.keyboard import KeyboardTeleop, KeyboardTeleopConfig from lerobot.teleoperators.so100_leader import SO100Leader, SO100LeaderConfig from lerobot.utils.control_utils import init_keyboard_listener from lerobot.utils.utils import log_say -from lerobot.utils.visualization_utils import _init_rerun +from lerobot.utils.visualization_utils import init_rerun NUM_EPISODES = 2 FPS = 30 @@ -69,7 +69,7 @@ keyboard.connect() # Initialize the keyboard listener and rerun visualization listener, events = init_keyboard_listener() -_init_rerun(session_name="lekiwi_record") +init_rerun(session_name="lekiwi_record") if not robot.is_connected or not leader_arm.is_connected or not keyboard.is_connected: raise ValueError("Robot or teleop is not connected!") diff --git a/examples/lekiwi/teleoperate.py b/examples/lekiwi/teleoperate.py index cde4000df..6b430df48 100644 --- a/examples/lekiwi/teleoperate.py +++ b/examples/lekiwi/teleoperate.py @@ -20,7 +20,7 @@ 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.visualization_utils import _init_rerun, log_rerun_data +from lerobot.utils.visualization_utils import init_rerun, log_rerun_data FPS = 30 @@ -41,7 +41,7 @@ leader_arm.connect() keyboard.connect() # Init rerun viewer -_init_rerun(session_name="lekiwi_teleop") +init_rerun(session_name="lekiwi_teleop") if not robot.is_connected or not leader_arm.is_connected or not keyboard.is_connected: raise ValueError("Robot or teleop is not connected!") diff --git a/examples/phone_to_so100/evaluate.py b/examples/phone_to_so100/evaluate.py index c7d6eb240..0d53f1177 100644 --- a/examples/phone_to_so100/evaluate.py +++ b/examples/phone_to_so100/evaluate.py @@ -43,7 +43,7 @@ from lerobot.robots.so100_follower.so100_follower import SO100Follower from lerobot.scripts.lerobot_record import record_loop from lerobot.utils.control_utils import init_keyboard_listener from lerobot.utils.utils import log_say -from lerobot.utils.visualization_utils import _init_rerun +from lerobot.utils.visualization_utils import init_rerun NUM_EPISODES = 5 FPS = 30 @@ -137,7 +137,7 @@ robot.connect() # Initialize the keyboard listener and rerun visualization listener, events = init_keyboard_listener() -_init_rerun(session_name="phone_so100_evaluate") +init_rerun(session_name="phone_so100_evaluate") if not robot.is_connected: raise ValueError("Robot is not connected!") diff --git a/examples/phone_to_so100/record.py b/examples/phone_to_so100/record.py index 6681017a0..bb2e2f5f7 100644 --- a/examples/phone_to_so100/record.py +++ b/examples/phone_to_so100/record.py @@ -41,7 +41,7 @@ from lerobot.teleoperators.phone.phone_processor import MapPhoneActionToRobotAct from lerobot.teleoperators.phone.teleop_phone import Phone from lerobot.utils.control_utils import init_keyboard_listener from lerobot.utils.utils import log_say -from lerobot.utils.visualization_utils import _init_rerun +from lerobot.utils.visualization_utils import init_rerun NUM_EPISODES = 2 FPS = 30 @@ -143,7 +143,7 @@ phone.connect() # Initialize the keyboard listener and rerun visualization listener, events = init_keyboard_listener() -_init_rerun(session_name="phone_so100_record") +init_rerun(session_name="phone_so100_record") if not robot.is_connected or not phone.is_connected: raise ValueError("Robot or teleop is not connected!") diff --git a/examples/phone_to_so100/teleoperate.py b/examples/phone_to_so100/teleoperate.py index eb5ed3526..6c49a8453 100644 --- a/examples/phone_to_so100/teleoperate.py +++ b/examples/phone_to_so100/teleoperate.py @@ -33,7 +33,7 @@ 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.visualization_utils import _init_rerun, log_rerun_data +from lerobot.utils.visualization_utils import init_rerun, log_rerun_data FPS = 30 @@ -87,7 +87,7 @@ robot.connect() teleop_device.connect() # Init rerun viewer -_init_rerun(session_name="phone_so100_teleop") +init_rerun(session_name="phone_so100_teleop") if not robot.is_connected or not teleop_device.is_connected: raise ValueError("Robot or teleop is not connected!") diff --git a/examples/so100_to_so100_EE/evaluate.py b/examples/so100_to_so100_EE/evaluate.py index f47a216d6..53a385442 100644 --- a/examples/so100_to_so100_EE/evaluate.py +++ b/examples/so100_to_so100_EE/evaluate.py @@ -43,7 +43,7 @@ from lerobot.robots.so100_follower.so100_follower import SO100Follower from lerobot.scripts.lerobot_record import record_loop from lerobot.utils.control_utils import init_keyboard_listener from lerobot.utils.utils import log_say -from lerobot.utils.visualization_utils import _init_rerun +from lerobot.utils.visualization_utils import init_rerun NUM_EPISODES = 5 FPS = 30 @@ -138,7 +138,7 @@ robot.connect() # Initialize the keyboard listener and rerun visualization listener, events = init_keyboard_listener() -_init_rerun(session_name="so100_so100_evaluate") +init_rerun(session_name="so100_so100_evaluate") if not robot.is_connected: raise ValueError("Robot is not connected!") diff --git a/examples/so100_to_so100_EE/record.py b/examples/so100_to_so100_EE/record.py index 60c96835f..6c38553e2 100644 --- a/examples/so100_to_so100_EE/record.py +++ b/examples/so100_to_so100_EE/record.py @@ -39,7 +39,7 @@ from lerobot.teleoperators.so100_leader.config_so100_leader import SO100LeaderCo from lerobot.teleoperators.so100_leader.so100_leader import SO100Leader from lerobot.utils.control_utils import init_keyboard_listener from lerobot.utils.utils import log_say -from lerobot.utils.visualization_utils import _init_rerun +from lerobot.utils.visualization_utils import init_rerun NUM_EPISODES = 2 FPS = 30 @@ -143,7 +143,7 @@ follower.connect() # Initialize the keyboard listener and rerun visualization listener, events = init_keyboard_listener() -_init_rerun(session_name="recording_phone") +init_rerun(session_name="recording_phone") if not leader.is_connected or not follower.is_connected: raise ValueError("Robot or teleop is not connected!") diff --git a/examples/so100_to_so100_EE/teleoperate.py b/examples/so100_to_so100_EE/teleoperate.py index ab54e7236..aa9755788 100644 --- a/examples/so100_to_so100_EE/teleoperate.py +++ b/examples/so100_to_so100_EE/teleoperate.py @@ -33,7 +33,7 @@ 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.visualization_utils import _init_rerun, log_rerun_data +from lerobot.utils.visualization_utils import init_rerun, log_rerun_data FPS = 30 @@ -95,7 +95,7 @@ follower.connect() leader.connect() # Init rerun viewer -_init_rerun(session_name="so100_so100_EE_teleop") +init_rerun(session_name="so100_so100_EE_teleop") print("Starting teleop loop...") while True: diff --git a/src/lerobot/scripts/lerobot_record.py b/src/lerobot/scripts/lerobot_record.py index d09b017e4..dd4984fab 100644 --- a/src/lerobot/scripts/lerobot_record.py +++ b/src/lerobot/scripts/lerobot_record.py @@ -122,7 +122,7 @@ from lerobot.utils.utils import ( init_logging, log_say, ) -from lerobot.utils.visualization_utils import _init_rerun, log_rerun_data +from lerobot.utils.visualization_utils import init_rerun, log_rerun_data @dataclass @@ -378,7 +378,7 @@ def record(cfg: RecordConfig) -> LeRobotDataset: init_logging() logging.info(pformat(asdict(cfg))) if cfg.display_data: - _init_rerun(session_name="recording") + init_rerun(session_name="recording") robot = make_robot_from_config(cfg.robot) teleop = make_teleoperator_from_config(cfg.teleop) if cfg.teleop is not None else None diff --git a/src/lerobot/scripts/lerobot_teleoperate.py b/src/lerobot/scripts/lerobot_teleoperate.py index 62c243e95..ab9a6361d 100644 --- a/src/lerobot/scripts/lerobot_teleoperate.py +++ b/src/lerobot/scripts/lerobot_teleoperate.py @@ -90,7 +90,7 @@ from lerobot.teleoperators import ( # noqa: F401 ) from lerobot.utils.robot_utils import busy_wait from lerobot.utils.utils import init_logging, move_cursor_up -from lerobot.utils.visualization_utils import _init_rerun, log_rerun_data +from lerobot.utils.visualization_utils import init_rerun, log_rerun_data @dataclass @@ -185,7 +185,7 @@ def teleoperate(cfg: TeleoperateConfig): init_logging() logging.info(pformat(asdict(cfg))) if cfg.display_data: - _init_rerun(session_name="teleoperation") + init_rerun(session_name="teleoperation") teleop = make_teleoperator_from_config(cfg.teleop) robot = make_robot_from_config(cfg.robot) diff --git a/src/lerobot/utils/robot_utils.py b/src/lerobot/utils/robot_utils.py index 8069b3662..42abcdda4 100644 --- a/src/lerobot/utils/robot_utils.py +++ b/src/lerobot/utils/robot_utils.py @@ -27,17 +27,3 @@ def busy_wait(seconds): # On Linux time.sleep is accurate if seconds > 0: time.sleep(seconds) - - -def safe_disconnect(func): - # TODO(aliberts): Allow to pass custom exceptions - # (e.g. ThreadServiceExit, KeyboardInterrupt, SystemExit, UnpluggedError, DynamixelCommError) - def wrapper(robot, *args, **kwargs): - try: - return func(robot, *args, **kwargs) - except Exception as e: - if robot.is_connected: - robot.disconnect() - raise e - - return wrapper diff --git a/src/lerobot/utils/train_utils.py b/src/lerobot/utils/train_utils.py index 08d1bcc9d..3ebe31971 100644 --- a/src/lerobot/utils/train_utils.py +++ b/src/lerobot/utils/train_utils.py @@ -13,10 +13,8 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -import logging from pathlib import Path -from termcolor import colored from torch.optim import Optimizer from torch.optim.lr_scheduler import LRScheduler @@ -36,10 +34,6 @@ from lerobot.utils.constants import ( from lerobot.utils.random_utils import load_rng_state, save_rng_state -def log_output_dir(out_dir): - logging.info(colored("Output dir:", "yellow", attrs=["bold"]) + f" {out_dir}") - - def get_step_identifier(step: int, total_steps: int) -> str: num_digits = max(6, len(str(total_steps))) return f"{step:0{num_digits}d}" diff --git a/src/lerobot/utils/utils.py b/src/lerobot/utils/utils.py index 107606fda..523a5e4d2 100644 --- a/src/lerobot/utils/utils.py +++ b/src/lerobot/utils/utils.py @@ -15,14 +15,13 @@ # limitations under the License. import logging import os -import os.path as osp import platform import select import subprocess import sys import time from copy import copy, deepcopy -from datetime import datetime, timezone +from datetime import datetime from pathlib import Path from statistics import mean @@ -30,12 +29,6 @@ import numpy as np import torch -def none_or_int(value): - if value == "None": - return None - return int(value) - - def inside_slurm(): """Check whether the python process was launched through slurm""" # TODO(rcadene): return False for interactive mode `--pty bash` @@ -165,36 +158,6 @@ def format_big_number(num, precision=0): return num -def _relative_path_between(path1: Path, path2: Path) -> Path: - """Returns path1 relative to path2.""" - path1 = path1.absolute() - path2 = path2.absolute() - try: - return path1.relative_to(path2) - except ValueError: # most likely because path1 is not a subpath of path2 - common_parts = Path(osp.commonpath([path1, path2])).parts - return Path( - "/".join([".."] * (len(path2.parts) - len(common_parts)) + list(path1.parts[len(common_parts) :])) - ) - - -def print_cuda_memory_usage(): - """Use this function to locate and debug memory leak.""" - import gc - - gc.collect() - # Also clear the cache if you want to fully release the memory - torch.cuda.empty_cache() - print(f"Current GPU Memory Allocated: {torch.cuda.memory_allocated(0) / 1024**2:.2f} MB") - print(f"Maximum GPU Memory Allocated: {torch.cuda.max_memory_allocated(0) / 1024**2:.2f} MB") - print(f"Current GPU Memory Reserved: {torch.cuda.memory_reserved(0) / 1024**2:.2f} MB") - print(f"Maximum GPU Memory Reserved: {torch.cuda.max_memory_reserved(0) / 1024**2:.2f} MB") - - -def capture_timestamp_utc(): - return datetime.now(timezone.utc) - - def say(text: str, blocking: bool = False): system = platform.system() diff --git a/src/lerobot/utils/visualization_utils.py b/src/lerobot/utils/visualization_utils.py index e6acc87de..7fc881f26 100644 --- a/src/lerobot/utils/visualization_utils.py +++ b/src/lerobot/utils/visualization_utils.py @@ -20,7 +20,7 @@ import numpy as np import rerun as rr -def _init_rerun(session_name: str = "lerobot_control_loop") -> None: +def init_rerun(session_name: str = "lerobot_control_loop") -> None: """Initializes the Rerun SDK for visualizing the control loop.""" batch_size = os.getenv("RERUN_FLUSH_NUM_BYTES", "8000") os.environ["RERUN_FLUSH_NUM_BYTES"] = batch_size