8c7cea1cb6
将插件猴子补丁改为通过 UF_LEROBOT_PATCH_TYPE 环境变量按需触发； SLAM 初始化增加就绪等待、重试及位姿置信度/时间戳校验； 新增 cameras_args 配置支持相机画面 resize； 修复遥操作 config parser 导入及多臂相机清理问题。
71 lines
3.6 KiB
Python
71 lines
3.6 KiB
Python
import os
|
||
|
||
# register plugin
|
||
import lerobot_robot_ufactory.cameras.umi_camera
|
||
import lerobot_robot_ufactory.robots.uf_robot
|
||
import lerobot_robot_ufactory.robots.uf_mock_robot
|
||
import lerobot_robot_ufactory.teleoperators.uf_mock_teleop
|
||
import lerobot_robot_ufactory.teleoperators.gello_teleop
|
||
import lerobot_robot_ufactory.teleoperators.pika_teleop
|
||
import lerobot_robot_ufactory.teleoperators.space_mouse
|
||
import lerobot_robot_ufactory.teleoperators.umi_teleop
|
||
|
||
def patch_lerobot_modules():
|
||
"""
|
||
Patch lerobot modules to use uFactory robot implementations.
|
||
This function replaces the default implementations of certain functions in the lerobot package
|
||
with the uFactory-specific implementations. It is intended to be called at the beginning of a script
|
||
that uses the lerobot package, before any other imports from lerobot are made.
|
||
"""
|
||
from lerobot_robot_ufactory.cameras.utils import make_cameras_from_configs as _uf_make_cameras_from_configs
|
||
from lerobot_robot_ufactory.robots.utils import make_robot_from_config as _uf_make_robot_from_config
|
||
from lerobot_robot_ufactory.teleoperators.utils import make_teleoperator_from_config as _uf_make_teleoperator_from_config
|
||
from lerobot_robot_ufactory.configs.parser import wrap as _uf_config_parser_wrap
|
||
import lerobot.cameras as _lerobot_cameras
|
||
import lerobot.robots as _lerobot_robot
|
||
import lerobot.teleoperators as _lerobot_teleoperators
|
||
import lerobot.cameras.utils as _lerobot_cameras_utils
|
||
import lerobot.robots.utils as _lerobot_robot_utils
|
||
import lerobot.teleoperators.utils as _lerobot_teleoperators_utils
|
||
import lerobot.configs.parser as _lerobot_configs_parser
|
||
# patch
|
||
_lerobot_cameras.make_cameras_from_configs = _uf_make_cameras_from_configs
|
||
_lerobot_robot.make_robot_from_config = _uf_make_robot_from_config
|
||
_lerobot_teleoperators.make_teleoperator_from_config = _uf_make_teleoperator_from_config
|
||
_lerobot_cameras_utils.make_cameras_from_configs = _uf_make_cameras_from_configs
|
||
_lerobot_robot_utils.make_robot_from_config = _uf_make_robot_from_config
|
||
_lerobot_teleoperators_utils.make_teleoperator_from_config = _uf_make_teleoperator_from_config
|
||
_lerobot_configs_parser.wrap = _uf_config_parser_wrap
|
||
|
||
_UF_LEROBOT_PATCH_TYPE = os.environ.get('UF_LEROBOT_PATCH_TYPE', '0')
|
||
if _UF_LEROBOT_PATCH_TYPE == '1':
|
||
print("Using UF_LEROBOT_PATCH_TYPE=1: patching lerobot modules for UFACTORY robot support.")
|
||
patch_lerobot_modules()
|
||
|
||
# # ── 重新装饰上游脚本中已用旧 wrap 装饰过的函数 ──
|
||
# # 因为 @parser.wrap() 在模块 import 时就执行了，
|
||
# # register_third_party_plugins() 触发本文件时 record() 已经是旧装饰器包好的。
|
||
# # 取出原始函数，用 _uf_config_wrap 重新包一次。
|
||
# import sys as _sys
|
||
# # 执行 python xxx.py 时模块名是 __main__，console_scripts 入口时是完整路径
|
||
# _SCRIPTS_TO_REWRAP = (
|
||
# "__main__",
|
||
# "lerobot.scripts.lerobot_record",
|
||
# "lerobot.scripts.lerobot_teleoperate",
|
||
# "lerobot.scripts.lerobot_eval",
|
||
# )
|
||
# for _mod_name in _SCRIPTS_TO_REWRAP:
|
||
# _mod = _sys.modules.get(_mod_name)
|
||
# if _mod is None:
|
||
# continue
|
||
# _record_fn = getattr(_mod, "record", getattr(_mod, "teleoperate", getattr(_mod, "eval_main", None)))
|
||
# if _record_fn is None:
|
||
# continue
|
||
# _original = getattr(_record_fn, "__wrapped__", None)
|
||
# if _original is not None:
|
||
# if _mod_name == "__main__":
|
||
# _sys.modules[_mod_name] = _uf_config_wrap()(_original)
|
||
# else:
|
||
# setattr(_mod, _record_fn.__name__, _uf_config_wrap()(_original))
|
||
# break # 找到并处理后就退出，避免重复
|