chore(processor): update comments in record.py

This commit is contained in:
Steven Palma
2025-09-16 18:04:44 +02:00
parent 43eb0e375f
commit d6ce7bd330
+7 -9
View File
@@ -205,7 +205,7 @@ class RecordConfig:
V V
[ robot.get_observation() ] ---> raw_obs [ robot.get_observation() ] ---> raw_obs
V V
[ robot_observation_processor ] ---> obs_transition [ robot_observation_processor ] ---> processed_obs
V V
.-----( ACTION LOGIC )------------------. .-----( ACTION LOGIC )------------------.
V V V V
@@ -216,7 +216,7 @@ class RecordConfig:
| V | V | V | V
| [teleop_action_processor] | | | [teleop_action_processor] | |
| | | | | | | |
'---> teleop_transition '---> policy_transition '---> processed_teleop_action '---> processed_policy_action
| | | |
'-------------------------.-------------' '-------------------------.-------------'
V V
@@ -224,7 +224,7 @@ class RecordConfig:
V V
[ robot.send_action() ] -- (Robot Executes) [ robot.send_action() ] -- (Robot Executes)
V V
( Transitions are merged & added to Dataset ) ( Save to Dataset )
V V
( Rerun Log / Loop Wait ) ( Rerun Log / Loop Wait )
""" """
@@ -253,7 +253,7 @@ def record_loop(
raise ValueError(f"The dataset fps should be equal to requested fps ({dataset.fps} != {fps}).") raise ValueError(f"The dataset fps should be equal to requested fps ({dataset.fps} != {fps}).")
teleop_arm = teleop_keyboard = None teleop_arm = teleop_keyboard = None
if isinstance(teleop, list): # For LeKiwi if isinstance(teleop, list):
teleop_keyboard = next((t for t in teleop if isinstance(t, KeyboardTeleop)), None) teleop_keyboard = next((t for t in teleop if isinstance(t, KeyboardTeleop)), None)
teleop_arm = next( teleop_arm = next(
( (
@@ -318,14 +318,13 @@ def record_loop(
f"{name}": float(action_values[i]) for i, name in enumerate(action_names) f"{name}": float(action_values[i]) for i, name in enumerate(action_names)
} }
elif isinstance(teleop, Teleoperator): elif policy is None and isinstance(teleop, Teleoperator):
act = teleop.get_action() act = teleop.get_action()
# Applies a pipeline to the raw teleop action, default is IdentityProcessor # Applies a pipeline to the raw teleop action, default is IdentityProcessor
# TODO(Steven): This assumes that the processor passed by the user should have identity_transition as to_output.
act_processed_teleop = teleop_action_processor(act) act_processed_teleop = teleop_action_processor(act)
elif isinstance(teleop, list): elif policy is None and isinstance(teleop, list):
arm_action = teleop_arm.get_action() arm_action = teleop_arm.get_action()
arm_action = {f"arm_{k}": v for k, v in arm_action.items()} arm_action = {f"arm_{k}": v for k, v in arm_action.items()}
keyboard_action = teleop_keyboard.get_action() keyboard_action = teleop_keyboard.get_action()
@@ -341,7 +340,6 @@ def record_loop(
continue continue
# Applies a pipeline to the action, default is IdentityProcessor # Applies a pipeline to the action, default is IdentityProcessor
# IMPORTANT: action_pipeline.to_output must return a dict suitable for robot.send_action()
if policy is not None and act_processed_policy is not None: if policy is not None and act_processed_policy is not None:
action_values = act_processed_policy action_values = act_processed_policy
robot_action_to_send = robot_action_processor(act_processed_policy) robot_action_to_send = robot_action_processor(act_processed_policy)
@@ -395,7 +393,7 @@ def record(cfg: RecordConfig) -> LeRobotDataset:
initial_features=create_initial_features(observation=robot.observation_features), initial_features=create_initial_features(observation=robot.observation_features),
use_videos=cfg.dataset.video, use_videos=cfg.dataset.video,
), ),
) # , **transition_features} )
if cfg.resume: if cfg.resume:
dataset = LeRobotDataset( dataset = LeRobotDataset(