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

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

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

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

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