From 357cbb11031957d1f6ba17b4bc2ccc0a6f016fdf Mon Sep 17 00:00:00 2001 From: blankey1337 <42594751+blankey1337@users.noreply.github.com> Date: Sun, 30 Nov 2025 07:44:23 -0800 Subject: [PATCH] refactor: align software structure with upstream guidelines - Remove custom `software/requirements.txt` to avoid conflict with upstream; moved to `examples/alohamini/requirements.txt` - Remove `lerobot_record.py` override; integrated lift logic into `LeKiwiClient` base action handler so standard script works - Move `dashboard` to `software/examples/alohamini/dashboard` to clean up root structure - Update `LeKiwiClient` to include `lift_axis.height_mm` in standard base action response This change ensures `software/src` can be used as a clean patch layer without modifying core LeRobot files. --- .../__pycache__/nav_service.cpython-310.pyc | Bin 0 -> 2892 bytes .../__pycache__/navigation.cpython-310.pyc | Bin 0 -> 2901 bytes .../{ => examples/alohamini}/dashboard/app.py | 0 .../alohamini}/dashboard/templates/index.html | 0 .../{ => examples/alohamini}/requirements.txt | 0 .../__pycache__/__init__.cpython-310.pyc | Bin 0 -> 362 bytes .../__pycache__/config_lekiwi.cpython-310.pyc | Bin 0 -> 2405 bytes .../lerobot/robots/alohamini/lekiwi_client.py | 44 +- .../src/lerobot/scripts/lerobot_record.py | 528 ------------------ 9 files changed, 7 insertions(+), 565 deletions(-) create mode 100644 software/examples/alohamini/__pycache__/nav_service.cpython-310.pyc create mode 100644 software/examples/alohamini/__pycache__/navigation.cpython-310.pyc rename software/{ => examples/alohamini}/dashboard/app.py (100%) rename software/{ => examples/alohamini}/dashboard/templates/index.html (100%) rename software/{ => examples/alohamini}/requirements.txt (100%) create mode 100644 software/src/lerobot/robots/alohamini/__pycache__/__init__.cpython-310.pyc create mode 100644 software/src/lerobot/robots/alohamini/__pycache__/config_lekiwi.cpython-310.pyc delete mode 100644 software/src/lerobot/scripts/lerobot_record.py diff --git a/software/examples/alohamini/__pycache__/nav_service.cpython-310.pyc b/software/examples/alohamini/__pycache__/nav_service.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d2ea6d8fca33f8b389a81f5dd6dc1fd2057491ad GIT binary patch literal 2892 zcmZuz&5s;M6|d^=>G{|XjO`5}O!!ckNbJNBQlteXv5hyNm`UWdvC>i0YEPBNy`JeF zSJ$rB!z>C~iyMCc1WPlB13AMb|4JRWaM}|GZk7oAUd@izc-^CV^Fswm-@@p&U>q?mx4< zaVMh(CL~wL18qXuAtb2E{(Y&9yL0P`iL1Rk>B&M|Nlj=u)~1HdovSx~apju5SjfIM z(bbQ7+t)AOyk@+QZ|(ft#N*Suq1r2oycn29e@`h{=yY6G()g+@N0o_%)Zl=?0A6JK zX{9s0S9Lt&RdRSYi6LoP-l@`joHUmBxKw&U^+y8il>U;b*P!In>)l(GR8?2p+0zKI z?#HU!1^cRdd05`dewi0}w<_=IhnbRHd6139Ls@k*+8fcH6<`r8o8w0&N>i-rG>yQ! zdIccjv9roo`9+LX=OPcBfYUGLlUL5f+vf6mpWLzh2q%Zw7^Ad@=aDb>9eC2|DO%}p z^(v+>x%xIHtJwm42c!veSObnFc`D+B$2ZPCesb!}+Z5yW?M)zpVwh>Z>p%)uaN+zO z*fPo~n{^;;&f7@ydN(_c33R3g2$Mi$a#R*C(Yy^WDlNrkCsY^ljH&5+N@jw>Ah>B! z6J&T4Ksoza`Bi`DzPK$dz> zQ7exewF2#JhNF5RqpYrOZ(IEkM?2m^h45ZhWm+q{XuRDrFJ2}l;#QaA@8RN!oH1{4 z@4u7tFO#zz_sN0Fk3^(WZ^hgOGKG6~**;e-%^@;M$`gkQut2cT&XjMkLvkI2pE|%ikBF6ua6FI>@} z88;fUCM0T2-8p393_JG3ZC1SIF}?EW*PI=CoK4pCs%TEVrC0UZY+bbI#%=Zqg7DD$ zjaT@K`5}L3^{&v~PXctMxA4}tV4K80pt3puYQMyPM0A3@lR6_5o}+|UT>#_T<%xf= zxi5!`D$-tneP6wYmO@#r)&NXh$v%~Bne98-I+L1{dlkI*}+ zp_F4pE0qJ|j;cZ8qDPou!P9X0LO`lFaPV!6>c;>9cMyj$k6aAQF=k-@X@#D1(TPyh z&pCnPecp0g{5+37cfM}@<7IWm7O)-+OC?oNm@4}7(5X@5pxj7P6Q}8@6njI$8)^Dr zFB>j)RGZFIEwY78{fMwj1W4BB#{v@=@x0%Syfw@RD%E+QvCoRp|*YU=g$7qu(GK4uw`6qE!49mDzq6Ri?NU`V~xj_45TC7Ydp54C;e*-)@)d7*PZv?eC^z*&}uXoWa_5q z_4>{MyTp^}bLiy@AQi7z!YTpED+^Rq_9r9}Z%Z1*K|x)~Kt0LMGg$o#OHGUXT3jBZ z=jgKbhIx^UCK^n-W0+v9vMWiZ)JWBSUg&I7rGvUPNK2jQe3(T}ZI{{Ll5$~Bc|?4W zx&tfYiB1OjFxTZxrVs=6&s*ysmYFKo<;97{{np=8qwTD(%k{H^(N_9SUgYcL=#u^- zRoQyBlMcs&tXxlN3F;1FA=w1C(Cj+D!DevgcB46B+$IIus=~ME1Fb5g6r5FConvW-F*I) z=3Bc>KA1FU$OB}@os4CwGfIx~B#6&_9w2&%AH~q}tz{kx+w%Deq8f@ogotd&Z82Rr z9NorHpDw%Ch!^N4v`4|1m(Wl=NH`dQJ8-j4w@0QjM?f0k%x(pPYw`r zm!jKHOiRjIp9MS!!l#Y455=KIQ#a7dr$H77#H9e*SkjUk0ev1rGI5e)&cykS$3_bP)wWSa0klUkHZxaU4{yt+8b@Ti0 zsQq##67zvSe%)ivMD`sBqwd4X`7qnqTk`dV-x}T57dGaDD;tq(TvP{xyvS0OjB{DH z(_(Wn*b{WBHZz@!M`czo=@FnQE&5p^^HRt6?aNq$p@*>Xoq?g)kvmFyDa+QfJ%^x2 z@8!`25*-2}N|r;e1q1+Wek3%91Go8gaf5N>_thcUPz4RQDeqMUHwBo2o7mlXFjitDUg|(BnbSOIMUW z>KP*Y32B~g09#4pxHm9tQj#8Lz#rn%$M1RX{S|#m&%I`#8-htHAFmd7L|Oray&ShO zk~qn}?V#fl8c_QLvn*lF?3n}bL=^-yd#H#GzP^U7ct<=iLxk6;wgx!e+Z_OCK=*8NRKT7IJw3 zasZp#q<^u|LoP33FYmfvwP7Qa?yLh4Z(rrIjk#Z8EWlU`W1Xra-70K$kHYCV=+4_7 z)KZe3?tZy*s=IrttG69upX4uf>&jnn<-K%b zya@-+x#z&eRn;+pI%DU$$O!pD^1m368(85IU?2lE#8N>m6O1BfCbPN7Z6S)pGgI1y zSim`NNmG1Nh&60Fy;pa7uUAH^F|ORW)z8z(K}!rg<;S2Q7KTVW<~ZX>`Tx-iCml)yKeZjeF@H z^r(4o+qmz%;=%VEJZ|G6PZ)TIE{iTH0i%=`?;q3h%DQ1^)HVJ*;G1<_)0+MQ7~xUz zz$N9#f54%=(KGxt`?zt+kt-LNS*c0e1o`b~c4l_9GxM8S7oARE!Lw|Coc!FgtiN#b z^5el|4Tky+2)8&%EsAeDC7Dg_%%Kj!xs$qCgElgcdS>pX&CI7h@D1*ztt_Ac@J-Pc zp=eC)Wov44pSPa4v?GG2T^>9m!oFegHV>azJQTs9b4U&e?HOr@cR|`cv_TRX`(AAy z@jht!$M$_rzO{yfBRH+KF|>7y$7MW8W2JOsnuwHxEL@K>A>;KTpC((ny}tM9W>Lmv zQshIT2X|*8U;pl9lOr28166yo7#C%|%#)&=%Y1?#difzRS%;xqAc+NsAsf#viBUcw z|NJm|poCN-KAxAUkfVFDcqk^N8hx1-+wnI^o{Us6E%#$7MoLabV1wTt8LCDxT4iW8 zF;1{l>?Hfi>g=&TSFvnjd{nG5xPT`@H(Vk^cxkUPJ9T_K^%@NI2#ByK;TE+YSi*Vg za>DIrHf;z`G`Vx7lDWJAGT+F&E0AH@5&`Ia;9C^3&0AktFb}}KBf7l(3+6Ne`#z!o zG$B}a=)kP%;wlJPPltEJZSbfHHjVCFyM62XTkOuY>-X8cyPNk9mPcIdjpB^OGGo&s zO^baYAq>+}w~(?~Axqk?*KOXqakE~ya^^yrY;Bi1FfX(8zP`v46^~QF%0lkW1uJp} z5(sAwiQ3Q$S^S7ekpgb^z$TVkqI^ztkBezMPfIq9C$R9bZq|vouG@1Z7#CYo2sK>L zK4UrLh%p^7mKA&spJF;>Y7WJ^F3%Y?PbQ_r9LpsnucA`6QPu_W2@HiyU@ba}BqUK|#f}`KI|Uj1 z#u_g88psJ`IH8VkMT6UrREK(ocY$xvrr|5N_=d0G(lR_ig|`6&p^hpL+$>aGbw|I0 z8#=_?vV&}wO)`G;2a3)ei;6K6)}EYJ4bL|i3*HbB0xAD5GDUThg@ zF)vxBj@|+vbcDH0i(T`JuwW4h_&~-u6sWm1Iv}^XLYP4r~tkE8%C(?tuM$K94N^MDz z6{TRwOuh?etitp@au}gg4q+75VW^8hu)%PL!;$g0cs=j{FN9wa4V?{5_BBzO#F}tT z4m5eBiPU7O$y}3%n(PkSv^gy#G)+$ZaXi^UPP?WlG)!sJw#2d#rcuzao>VGd&4l1= zKBGa!@nS!x%`|~#0>YMIP^R78c~ZWfU6X7qAnjjC<7~|1t3RS+h#1?p$?*?)GrYeK z-Z&k>)fcPf0&dnWMJfP@JL0i2HGE|Hz+eU~FUenlE@1iu?8lf^8^h>z>|(Gh*q)%+ z0N6prXGks~Ic^+P%lI5+Um!UP*_T1F2}3Oc2`t|QWP_(3S^9TnF|;Eh(PB7|EAXQ2 z<9k^sX@kVHDtf-Y=vM2lxeB!lm$a%5fF9DzubUKbVNdpI%egeA&nltBzc19QHjzGl z<`NPU0Q8RbrPxYT33rE@kInxC)8tGwnggHZE2Ow;*jMqkH`w*N5s7SY* literal 0 HcmV?d00001 diff --git a/software/src/lerobot/robots/alohamini/lekiwi_client.py b/software/src/lerobot/robots/alohamini/lekiwi_client.py index 94f8fa8..1fdfd22 100644 --- a/software/src/lerobot/robots/alohamini/lekiwi_client.py +++ b/software/src/lerobot/robots/alohamini/lekiwi_client.py @@ -324,56 +324,26 @@ def _from_keyboard_to_base_action(self, pressed_keys: np.ndarray): if self.teleop_keys["rotate_right"] in pressed_keys: theta_cmd -= theta_speed - - return { - "x.vel": x_cmd, - "y.vel": y_cmd, - "theta.vel": theta_cmd, - } - - # lift_axis.vel - # def _from_keyboard_to_lift_action(self, pressed_keys: np.ndarray): - # LIFT_VEL = 1000 # 觉得慢/快就改 - # up_pressed = self.teleop_keys.get("lift_up", "u") in pressed_keys - # dn_pressed = self.teleop_keys.get("lift_down", "j") in pressed_keys - - # if up_pressed and not dn_pressed: - # v = +LIFT_VEL - # elif dn_pressed and not up_pressed: - # v = -LIFT_VEL - # else: - # v = 0.0 - # return {"lift_axis.vel": int(v)} - - - # lift_axis.height_mm - def _from_keyboard_to_lift_action(self, pressed_keys: np.ndarray): + # Lift Control up_pressed = self.teleop_keys.get("lift_up", "u") in pressed_keys dn_pressed = self.teleop_keys.get("lift_down", "j") in pressed_keys # Read the last height (mm) reported by the Host h_now = float(self.last_remote_state.get("lift_axis.height_mm", 0.0)) - #print(f"h_now:{h_now}") - if not (up_pressed or dn_pressed): - # If neither 'u' nor 'j' is pressed, reuse the previous value to avoid empty input - #return {"lift_axis.height_mm": h_now} - return {"lift_axis.height_mm": h_now} - - # Increment on each key press if up_pressed and not dn_pressed: h_target = h_now + LiftAxisConfig.step_mm elif dn_pressed and not up_pressed: h_target = h_now - LiftAxisConfig.step_mm else: h_target = h_now - #print(f"h_target:{h_target}") - - # Send "target height (mm)" directly - return {"lift_axis.height_mm": h_target} - - + return { + "x.vel": x_cmd, + "y.vel": y_cmd, + "theta.vel": theta_cmd, + "lift_axis.height_mm": h_target, + } def configure(self): pass diff --git a/software/src/lerobot/scripts/lerobot_record.py b/software/src/lerobot/scripts/lerobot_record.py deleted file mode 100644 index 4d46af9..0000000 --- a/software/src/lerobot/scripts/lerobot_record.py +++ /dev/null @@ -1,528 +0,0 @@ -# Copyright 2024 The HuggingFace Inc. team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# 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. - -""" -Records a dataset. Actions for the robot can be either generated by teleoperation or by a policy. - -Example: - -```shell -lerobot-record \ - --robot.type=so100_follower \ - --robot.port=/dev/tty.usbmodem58760431541 \ - --robot.cameras="{laptop: {type: opencv, index_or_path: 0, width: 640, height: 480, fps: 30}}" \ - --robot.id=black \ - --dataset.repo_id=/ \ - --dataset.num_episodes=2 \ - --dataset.single_task="Grab the cube" \ - --display_data=true - # <- Teleop optional if you want to teleoperate to record or in between episodes with a policy \ - # --teleop.type=so100_leader \ - # --teleop.port=/dev/tty.usbmodem58760431551 \ - # --teleop.id=blue \ - # <- Policy optional if you want to record with a policy \ - # --policy.path=${HF_USER}/my_policy \ -``` - -Example recording with bimanual so100: -```shell -lerobot-record \ - --robot.type=bi_so100_follower \ - --robot.left_arm_port=/dev/tty.usbmodem5A460851411 \ - --robot.right_arm_port=/dev/tty.usbmodem5A460812391 \ - --robot.id=bimanual_follower \ - --robot.cameras='{ - left: {"type": "opencv", "index_or_path": 0, "width": 640, "height": 480, "fps": 30}, - top: {"type": "opencv", "index_or_path": 1, "width": 640, "height": 480, "fps": 30}, - right: {"type": "opencv", "index_or_path": 2, "width": 640, "height": 480, "fps": 30} - }' \ - --teleop.type=bi_so100_leader \ - --teleop.left_arm_port=/dev/tty.usbmodem5A460828611 \ - --teleop.right_arm_port=/dev/tty.usbmodem5A460826981 \ - --teleop.id=bimanual_leader \ - --display_data=true \ - --dataset.repo_id=${HF_USER}/bimanual-so100-handover-cube \ - --dataset.num_episodes=25 \ - --dataset.single_task="Grab and handover the red cube to the other arm" -``` -""" - -import logging -import time -from dataclasses import asdict, dataclass, field -from pathlib import Path -from pprint import pformat -from typing import Any - -from lerobot.cameras import ( # noqa: F401 - CameraConfig, # noqa: F401 -) -from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig # noqa: F401 -from lerobot.cameras.realsense.configuration_realsense import RealSenseCameraConfig # noqa: F401 -from lerobot.configs import parser -from lerobot.configs.policies import PreTrainedConfig -from lerobot.datasets.image_writer import safe_stop_image_writer -from lerobot.datasets.lerobot_dataset import LeRobotDataset -from lerobot.datasets.pipeline_features import aggregate_pipeline_dataset_features, create_initial_features -from lerobot.datasets.utils import build_dataset_frame, combine_feature_dicts -from lerobot.datasets.video_utils import VideoEncodingManager -from lerobot.policies.factory import make_policy, make_pre_post_processors -from lerobot.policies.pretrained import PreTrainedPolicy -from lerobot.policies.utils import make_robot_action -from lerobot.processor import ( - PolicyAction, - PolicyProcessorPipeline, - RobotAction, - RobotObservation, - RobotProcessorPipeline, - make_default_processors, -) -from lerobot.processor.rename_processor import rename_stats -from lerobot.robots import ( # noqa: F401 - Robot, - RobotConfig, - bi_so100_follower, - hope_jr, - koch_follower, - make_robot_from_config, - so100_follower, - so101_follower, -) -from lerobot.teleoperators import ( # noqa: F401 - Teleoperator, - TeleoperatorConfig, - bi_so100_leader, - homunculus, - koch_leader, - make_teleoperator_from_config, - so100_leader, - so101_leader, -) -from lerobot.teleoperators.keyboard.teleop_keyboard import KeyboardTeleop -from lerobot.utils.constants import ACTION, OBS_STR -from lerobot.utils.control_utils import ( - init_keyboard_listener, - is_headless, - predict_action, - sanity_check_dataset_name, - 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.utils import ( - get_safe_torch_device, - init_logging, - log_say, -) -from lerobot.utils.visualization_utils import init_rerun, log_rerun_data - - -@dataclass -class DatasetRecordConfig: - # Dataset identifier. By convention it should match '{hf_username}/{dataset_name}' (e.g. `lerobot/test`). - repo_id: str - # A short but accurate description of the task performed during the recording (e.g. "Pick the Lego block and drop it in the box on the right.") - single_task: str - # Root directory where the dataset will be stored (e.g. 'dataset/path'). - root: str | Path | None = None - # Limit the frames per second. - fps: int = 30 - # Number of seconds for data recording for each episode. - episode_time_s: int | float = 60 - # Number of seconds for resetting the environment after each episode. - reset_time_s: int | float = 60 - # Number of episodes to record. - num_episodes: int = 50 - # Encode frames in the dataset into video - video: bool = True - # Upload dataset to Hugging Face hub. - push_to_hub: bool = True - # Upload on private repository on the Hugging Face hub. - private: bool = False - # Add tags to your dataset on the hub. - tags: list[str] | None = None - # Number of subprocesses handling the saving of frames as PNG. Set to 0 to use threads only; - # set to ≥1 to use subprocesses, each using threads to write images. The best number of processes - # and threads depends on your system. We recommend 4 threads per camera with 0 processes. - # If fps is unstable, adjust the thread count. If still unstable, try using 1 or more subprocesses. - num_image_writer_processes: int = 0 - # Number of threads writing the frames as png images on disk, per camera. - # Too many threads might cause unstable teleoperation fps due to main thread being blocked. - # Not enough threads might cause low camera fps. - num_image_writer_threads_per_camera: int = 4 - # Number of episodes to record before batch encoding videos - # Set to 1 for immediate encoding (default behavior), or higher for batched encoding - video_encoding_batch_size: int = 1 - # Rename map for the observation to override the image and state keys - rename_map: dict[str, str] = field(default_factory=dict) - - def __post_init__(self): - if self.single_task is None: - raise ValueError("You need to provide a task as argument in `single_task`.") - - -@dataclass -class RecordConfig: - robot: RobotConfig - dataset: DatasetRecordConfig - # Whether to control the robot with a teleoperator - teleop: TeleoperatorConfig | None = None - # Whether to control the robot with a policy - policy: PreTrainedConfig | None = None - # Display all cameras on screen - display_data: bool = False - # Use vocal synthesis to read events. - play_sounds: bool = True - # Resume recording on an existing dataset. - resume: bool = False - - def __post_init__(self): - # HACK: We parse again the cli args here to get the pretrained path if there was one. - policy_path = parser.get_path_arg("policy") - if policy_path: - cli_overrides = parser.get_cli_overrides("policy") - self.policy = PreTrainedConfig.from_pretrained(policy_path, cli_overrides=cli_overrides) - self.policy.pretrained_path = policy_path - - if self.teleop is None and self.policy is None: - raise ValueError("Choose a policy, a teleoperator or both to control the robot") - - @classmethod - def __get_path_fields__(cls) -> list[str]: - """This enables the parser to load config from the policy using `--policy.path=local/dir`""" - return ["policy"] - - -""" --------------- record_loop() data flow -------------------------- - [ Robot ] - V - [ robot.get_observation() ] ---> raw_obs - V - [ robot_observation_processor ] ---> processed_obs - V - .-----( ACTION LOGIC )------------------. - V V - [ From Teleoperator ] [ From Policy ] - | | - | [teleop.get_action] -> raw_action | [predict_action] - | | | | - | V | V - | [teleop_action_processor] | | - | | | | - '---> processed_teleop_action '---> processed_policy_action - | | - '-------------------------.-------------' - V - [ robot_action_processor ] --> robot_action_to_send - V - [ robot.send_action() ] -- (Robot Executes) - V - ( Save to Dataset ) - V - ( Rerun Log / Loop Wait ) -""" - - -@safe_stop_image_writer -def record_loop( - robot: Robot, - events: dict, - fps: int, - teleop_action_processor: RobotProcessorPipeline[ - tuple[RobotAction, RobotObservation], RobotAction - ], # runs after teleop - robot_action_processor: RobotProcessorPipeline[ - tuple[RobotAction, RobotObservation], RobotAction - ], # runs before robot - robot_observation_processor: RobotProcessorPipeline[ - RobotObservation, RobotObservation - ], # runs after robot - dataset: LeRobotDataset | None = None, - teleop: Teleoperator | list[Teleoperator] | None = None, - policy: PreTrainedPolicy | None = None, - preprocessor: PolicyProcessorPipeline[dict[str, Any], dict[str, Any]] | None = None, - postprocessor: PolicyProcessorPipeline[PolicyAction, PolicyAction] | None = None, - control_time_s: int | None = None, - single_task: str | None = None, - display_data: bool = False, -): - if dataset is not None and dataset.fps != fps: - raise ValueError(f"The dataset fps should be equal to requested fps ({dataset.fps} != {fps}).") - - teleop_arm = teleop_keyboard = None - if isinstance(teleop, list): - teleop_keyboard = next((t for t in teleop if isinstance(t, KeyboardTeleop)), None) - teleop_arm = next( - ( - t - for t in teleop - if isinstance( - t, - (so100_leader.SO100Leader | so101_leader.SO101Leader | koch_leader.KochLeader| bi_so100_leader.BiSO100Leader), - ) - ), - None, - ) - - if not (teleop_arm and teleop_keyboard and len(teleop) == 2 and robot.name == "lekiwi_client"): - raise ValueError( - f"For multi-teleop, the list must contain exactly one KeyboardTeleop and one arm teleoperator. " - f"Currently only supported for LeKiwi robot.\n" - f"Got values:\n" - f" robot.name = '{getattr(robot, 'name', None)}'\n" - f" len(teleop) = {len(teleop)}\n" - f" teleop_arm = {bool(teleop_arm)}\n" - f" teleop_keyboard = {bool(teleop_keyboard)}\n" - f" teleop types = {[type(t).__name__ for t in teleop]}" - ) - - # Reset policy and processor if they are provided - if policy is not None and preprocessor is not None and postprocessor is not None: - policy.reset() - preprocessor.reset() - postprocessor.reset() - - timestamp = 0 - start_episode_t = time.perf_counter() - while timestamp < control_time_s: - start_loop_t = time.perf_counter() - - if events["exit_early"]: - events["exit_early"] = False - break - - # Get robot observation - obs = robot.get_observation() - - # Applies a pipeline to the raw robot observation, default is IdentityProcessor - obs_processed = robot_observation_processor(obs) - - if policy is not None or dataset is not None: - observation_frame = build_dataset_frame(dataset.features, obs_processed, prefix=OBS_STR) - - # Get action from either policy or teleop - if policy is not None and preprocessor is not None and postprocessor is not None: - action_values = predict_action( - observation=observation_frame, - policy=policy, - device=get_safe_torch_device(policy.config.device), - preprocessor=preprocessor, - postprocessor=postprocessor, - use_amp=policy.config.use_amp, - task=single_task, - robot_type=robot.robot_type, - ) - - act_processed_policy: RobotAction = make_robot_action(action_values, dataset.features) - - elif policy is None and isinstance(teleop, Teleoperator): - act = teleop.get_action() - - # Applies a pipeline to the raw teleop action, default is IdentityProcessor - act_processed_teleop = teleop_action_processor((act, obs)) - - elif policy is None and isinstance(teleop, list): - arm_action = teleop_arm.get_action() - arm_action = {f"arm_{k}": v for k, v in arm_action.items()} - keyboard_action = teleop_keyboard.get_action() - base_action = robot._from_keyboard_to_base_action(keyboard_action) - lift_action = robot._from_keyboard_to_lift_action(keyboard_action) #AlohaMini addition - act = {**arm_action, **base_action, **lift_action} if len(base_action) > 0 else arm_action #AlohaMini fix - act_processed_teleop = teleop_action_processor((act, obs)) - else: - logging.info( - "No policy or teleoperator provided, skipping action generation." - "This is likely to happen when resetting the environment without a teleop device." - "The robot won't be at its rest position at the start of the next episode." - ) - continue - - # Applies a pipeline to the action, default is IdentityProcessor - if policy is not None and act_processed_policy is not None: - action_values = act_processed_policy - robot_action_to_send = robot_action_processor((act_processed_policy, obs)) - else: - action_values = act_processed_teleop - robot_action_to_send = robot_action_processor((act_processed_teleop, obs)) - - # Send action to robot - # Action can eventually be clipped using `max_relative_target`, - # so action actually sent is saved in the dataset. action = postprocessor.process(action) - # TODO(steven, pepijn, adil): we should use a pipeline step to clip the action, so the sent action is the action that we input to the robot. - _sent_action = robot.send_action(robot_action_to_send) - - # Write to dataset - if dataset is not None: - action_frame = build_dataset_frame(dataset.features, action_values, prefix=ACTION) - frame = {**observation_frame, **action_frame, "task": single_task} - dataset.add_frame(frame) - - if display_data: - log_rerun_data(observation=obs_processed, action=action_values) - - dt_s = time.perf_counter() - start_loop_t - busy_wait(1 / fps - dt_s) - - timestamp = time.perf_counter() - start_episode_t - - -@parser.wrap() -def record(cfg: RecordConfig) -> LeRobotDataset: - init_logging() - logging.info(pformat(asdict(cfg))) - if cfg.display_data: - 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 - - teleop_action_processor, robot_action_processor, robot_observation_processor = make_default_processors() - - dataset_features = combine_feature_dicts( - aggregate_pipeline_dataset_features( - pipeline=teleop_action_processor, - initial_features=create_initial_features( - action=robot.action_features - ), # TODO(steven, pepijn): in future this should be come from teleop or policy - use_videos=cfg.dataset.video, - ), - aggregate_pipeline_dataset_features( - pipeline=robot_observation_processor, - initial_features=create_initial_features(observation=robot.observation_features), - use_videos=cfg.dataset.video, - ), - ) - - if cfg.resume: - dataset = LeRobotDataset( - cfg.dataset.repo_id, - root=cfg.dataset.root, - batch_encoding_size=cfg.dataset.video_encoding_batch_size, - ) - - if hasattr(robot, "cameras") and len(robot.cameras) > 0: - dataset.start_image_writer( - num_processes=cfg.dataset.num_image_writer_processes, - num_threads=cfg.dataset.num_image_writer_threads_per_camera * len(robot.cameras), - ) - sanity_check_dataset_robot_compatibility(dataset, robot, cfg.dataset.fps, dataset_features) - else: - # Create empty dataset or load existing saved episodes - sanity_check_dataset_name(cfg.dataset.repo_id, cfg.policy) - dataset = LeRobotDataset.create( - cfg.dataset.repo_id, - cfg.dataset.fps, - root=cfg.dataset.root, - robot_type=robot.name, - features=dataset_features, - use_videos=cfg.dataset.video, - image_writer_processes=cfg.dataset.num_image_writer_processes, - image_writer_threads=cfg.dataset.num_image_writer_threads_per_camera * len(robot.cameras), - batch_encoding_size=cfg.dataset.video_encoding_batch_size, - ) - - # Load pretrained policy - policy = None if cfg.policy is None else make_policy(cfg.policy, ds_meta=dataset.meta) - preprocessor = None - postprocessor = None - if cfg.policy is not None: - preprocessor, postprocessor = make_pre_post_processors( - policy_cfg=cfg.policy, - pretrained_path=cfg.policy.pretrained_path, - dataset_stats=rename_stats(dataset.meta.stats, cfg.dataset.rename_map), - preprocessor_overrides={ - "device_processor": {"device": cfg.policy.device}, - "rename_observations_processor": {"rename_map": cfg.dataset.rename_map}, - }, - ) - - robot.connect() - if teleop is not None: - teleop.connect() - - listener, events = init_keyboard_listener() - - with VideoEncodingManager(dataset): - recorded_episodes = 0 - while recorded_episodes < cfg.dataset.num_episodes and not events["stop_recording"]: - log_say(f"Recording episode {dataset.num_episodes}", cfg.play_sounds) - record_loop( - robot=robot, - events=events, - fps=cfg.dataset.fps, - teleop_action_processor=teleop_action_processor, - robot_action_processor=robot_action_processor, - robot_observation_processor=robot_observation_processor, - teleop=teleop, - policy=policy, - preprocessor=preprocessor, - postprocessor=postprocessor, - dataset=dataset, - control_time_s=cfg.dataset.episode_time_s, - single_task=cfg.dataset.single_task, - display_data=cfg.display_data, - ) - - # Execute a few seconds without recording to give time to manually reset the environment - # Skip reset for the last episode to be recorded - if not events["stop_recording"] and ( - (recorded_episodes < cfg.dataset.num_episodes - 1) or events["rerecord_episode"] - ): - log_say("Reset the environment", cfg.play_sounds) - record_loop( - robot=robot, - events=events, - fps=cfg.dataset.fps, - teleop_action_processor=teleop_action_processor, - robot_action_processor=robot_action_processor, - robot_observation_processor=robot_observation_processor, - teleop=teleop, - control_time_s=cfg.dataset.reset_time_s, - single_task=cfg.dataset.single_task, - display_data=cfg.display_data, - ) - - if events["rerecord_episode"]: - log_say("Re-record episode", cfg.play_sounds) - events["rerecord_episode"] = False - events["exit_early"] = False - dataset.clear_episode_buffer() - continue - - dataset.save_episode() - recorded_episodes += 1 - - log_say("Stop recording", cfg.play_sounds, blocking=True) - - robot.disconnect() - if teleop is not None: - teleop.disconnect() - - if not is_headless() and listener is not None: - listener.stop() - - if cfg.dataset.push_to_hub: - dataset.push_to_hub(tags=cfg.dataset.tags, private=cfg.dataset.private) - - log_say("Exiting", cfg.play_sounds) - return dataset - - -def main(): - register_third_party_devices() - record() - - -if __name__ == "__main__": - main()