wangzixiang
/
lerobot
mirror of https://github.com/huggingface/lerobot.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Update teleop_step for conditional joint state recording

This commit is contained in:
shantanu-tr 2025-04-04 11:33:57 -05:00
parent 5db796f3ff
commit ab53095728
2 changed files with 6 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
lerobot/common/robot_devices/robots/manipulator.py
View File

@ -462,7 +462,7 @@ class ManipulatorRobot:
self.follower_arms[name].write("Acceleration", 254)
def teleop_step(
self, record_data=False
self, record_data=False, record_joint_angles=True,
) -> None | tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]:
if not self.is_connected:
raise RobotDeviceNotConnectedError(
@ -537,6 +537,7 @@ class ManipulatorRobot:
# Populate output dictionaries
obs_dict, action_dict = {}, {}
if record_joint_angles:
obs_dict["observation.state"] = state
action_dict["action"] = action
for name in self.cameras:

3
lerobot/common/robot_devices/robots/trossen_ai_mobile.py
View File

@ -201,7 +201,7 @@ class TrossenAIMobile():
}
def teleop_step(
self, record_data=False
self, record_data=False, record_joint_angles=True,
) -> None | tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]:
if not self.is_connected:
@ -280,6 +280,7 @@ class TrossenAIMobile():
# Populate output dictionaries
obs_dict, action_dict = {}, {}
if record_joint_angles:
obs_dict["observation.state"] = state
action_dict["action"] = action
for name in self.cameras: