Skip to content

Commit c02bfe9

Browse files
juelgjuelg
authored
Merge pull request #297 from RobotControlStack/juelg/teleop-features
feat(recording): dataset inspection and conversion
2 parents 779d06b + a2a362d commit c02bfe9

6 files changed

Lines changed: 1045 additions & 4 deletions

File tree

‎examples/teleop/inspect_parquet_duckdb.ipynb‎

Lines changed: 505 additions & 0 deletions
Large diffs are not rendered by default.

‎examples/teleop/requirements.txt‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,6 @@
11
# simpub @ git+https://github.com/intuitive-robots/SimPublisher.git@a0ebdbfba86f58aaa6360eb3b9da7b74a9c9b1dd # doesnt work
22
dynamixel_sdk
33
pynput
4-
scipy
4+
scipy
5+
jupyter
6+
mutplotlib

‎python/rcs/__init__.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,11 @@ class RobotMetaConfig:
125125
common.GripperType("Robotiq2F85"): "assets/grippers/robotiq_2f85/robotiq_2f85.xml",
126126
}
127127

128+
GRIPPER_OFFSETS: dict[common.GripperType, common.Pose] = {
129+
common.GripperType.FrankaHand: common.Pose(pose_matrix=common.FrankaHandTCPOffset()),
130+
common.GripperType("Robotiq2F85"): common.Pose(translation=np.array([0.0, 0.0, 0.1628])),
131+
}
132+
128133
SCENE_PATHS: dict[str, str] = {"empty_world": "assets/scenes/empty_world/scene.xml"}
129134

130135
OBJECT_PATHS: dict[str, str] = {

‎python/rcs/__main__.py‎

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,21 @@
33

44
import typer
55
from rcs.envs.storage_wrapper import StorageWrapper
6+
from rcs.lerobot_joint_converter import (
7+
DEFAULT_CAMERAS,
8+
DEFAULT_DATASET_PATHS,
9+
DEFAULT_FPS,
10+
DEFAULT_GRIPPER_TYPE,
11+
DEFAULT_HF_DATA_DIR,
12+
DEFAULT_IMAGE_BATCH_SIZE,
13+
DEFAULT_JOINTS,
14+
DEFAULT_PER_ROBOT_ARM_DIM,
15+
DEFAULT_REPO_ID,
16+
DEFAULT_ROBOT_KEYS,
17+
DEFAULT_ROBOT_TYPE,
18+
camera_specs_to_configs,
19+
run_conversion,
20+
)
621
from rcs.sim.replayer import replay as replay_dataset
722

823
app = typer.Typer()
@@ -67,5 +82,80 @@ def replay(
6782
)
6883

6984

85+
@app.command("lerobot-convert")
86+
def lerobot_convert(
87+
output: Annotated[
88+
Path,
89+
typer.Argument(
90+
help="Output directory for the LeRobot dataset. Example: --output ./data_lerobot",
91+
),
92+
] = Path(DEFAULT_HF_DATA_DIR),
93+
dataset_paths: Annotated[
94+
list[str] | None,
95+
typer.Option(
96+
"--dataset-path",
97+
help="Input parquet path or glob. Repeat for multiple datasets. Example: --dataset-path /data/session1 --dataset-path /data/session2",
98+
),
99+
] = None,
100+
repo_id: Annotated[
101+
str, typer.Option(help="LeRobot repo id metadata. Example: --repo-id myorg/grasp_v2")
102+
] = DEFAULT_REPO_ID,
103+
robot_type: Annotated[
104+
str, typer.Option(help="Robot type for metadata and IK model lookup. Example: --robot-type fr3")
105+
] = DEFAULT_ROBOT_TYPE,
106+
fps: Annotated[int, typer.Option(help="Dataset frames per second. Example: --fps 30")] = DEFAULT_FPS,
107+
robot_keys: Annotated[
108+
list[str] | None,
109+
typer.Option(
110+
"--robot-key",
111+
help="Robot keys to concatenate. Repeat for multiple robots. Example: --robot-key left --robot-key right",
112+
),
113+
] = None,
114+
joints: Annotated[
115+
bool, typer.Option(help="Whether absolute_action is already in joint space. Example: --joints")
116+
] = DEFAULT_JOINTS,
117+
gripper_type: Annotated[
118+
str, typer.Option(help="Gripper type used to derive TCP offset. Example: --gripper-type Robotiq2F85")
119+
] = DEFAULT_GRIPPER_TYPE,
120+
camera_specs: Annotated[
121+
list[str] | None,
122+
typer.Option(
123+
"--camera",
124+
help=(
125+
"Camera spec as name[:source_name][@HEIGHTxWIDTH]. Repeat for multiple cameras. "
126+
"The name becomes the LeRobot output key (observation.images.<name>). "
127+
"The optional source_name is the key in the source parquet (obs.frames.<source_name>.rgb.data); "
128+
"if omitted, the image_ prefix is stripped from name to derive it. "
129+
"Example: --camera head@256x256 --camera image_left_wrist:left_wrist@256x256"
130+
),
131+
),
132+
] = None,
133+
image_batch_size: Annotated[
134+
int, typer.Option(help="Batch size for image decoding. Example: --image-batch-size 32")
135+
] = DEFAULT_IMAGE_BATCH_SIZE,
136+
per_robot_arm_dim: Annotated[
137+
int, typer.Option(help="Per-robot arm joint/action dimension without gripper. Example: --per-robot-arm-dim 7")
138+
] = DEFAULT_PER_ROBOT_ARM_DIM,
139+
success: Annotated[bool, typer.Option(help="Only include successful episodes. Example: --success")] = True,
140+
n: Annotated[int, typer.Option(help="Maximum number of episodes to convert. -1 means all. Example: --n 50")] = -1,
141+
):
142+
cameras = camera_specs_to_configs(camera_specs) if camera_specs is not None else list(DEFAULT_CAMERAS)
143+
run_conversion(
144+
root=output,
145+
dataset_paths=dataset_paths or list(DEFAULT_DATASET_PATHS),
146+
repo_id=repo_id,
147+
robot_type=robot_type,
148+
fps=fps,
149+
robot_keys=robot_keys or list(DEFAULT_ROBOT_KEYS),
150+
joints=joints,
151+
gripper_type=gripper_type,
152+
cameras=cameras,
153+
image_batch_size=image_batch_size,
154+
per_robot_arm_dim=per_robot_arm_dim,
155+
success=success,
156+
n=n,
157+
)
158+
159+
70160
if __name__ == "__main__":
71161
app()

‎python/rcs/envs/configs.py‎

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,13 @@
2121
from rcs.envs.tasks import PickTaskConfig
2222

2323
import rcs
24-
from rcs import CAMERA_PATHS, DEFAULT_TRANSFORMS, OBJECT_PATHS, SCENE_PATHS
24+
from rcs import (
25+
CAMERA_PATHS,
26+
DEFAULT_TRANSFORMS,
27+
GRIPPER_OFFSETS,
28+
OBJECT_PATHS,
29+
SCENE_PATHS,
30+
)
2531

2632

2733
class EmptyWorldFR3(SimEnvCreator):
@@ -33,7 +39,7 @@ def config(self) -> SimEnvCreatorConfig:
3339
q_home[-1] = np.pi / 4
3440
robot_cfg = SimRobotConfig(
3541
robot_type=RobotType.FR3,
36-
tcp_offset=rcs.common.Pose(pose_matrix=FrankaHandTCPOffset()),
42+
tcp_offset=GRIPPER_OFFSETS[rcs.common.GripperType.FrankaHand],
3743
attachment_site=rcs.ROBOTS[RobotType.FR3].attachment_site,
3844
kinematic_model_path=rcs.ROBOTS[RobotType.FR3].mjcf_model_path,
3945
joint_rotational_tolerance=0.05 * (np.pi / 180.0),
@@ -178,6 +184,7 @@ class EmptyWorldFR3Duo(SimEnvCreator):
178184

179185
def config(self) -> SimEnvCreatorConfig:
180186
robot_cfg = SimRobotConfig(
187+
tcp_offset=GRIPPER_OFFSETS[rcs.common.GripperType("Robotiq2F85")],
181188
robot_type=RobotType.FR3,
182189
attachment_site=rcs.ROBOTS[RobotType.FR3].attachment_site,
183190
kinematic_model_path=rcs.ROBOTS[RobotType.FR3].mjcf_model_path,
@@ -357,7 +364,7 @@ def config(self) -> SimEnvCreatorConfig:
357364
lead_robot_name = self.lead_robot_name(cfg)
358365

359366
robot_cfg = cfg.robot_cfgs[lead_robot_name]
360-
robot_cfg.tcp_offset = rcs.common.Pose()
367+
robot_cfg.tcp_offset = GRIPPER_OFFSETS[rcs.common.GripperType("Robotiq2F85")]
361368
robot_cfg.attachment_site = rcs.ROBOTS[rt].attachment_site
362369
robot_cfg.kinematic_model_path = rcs.ROBOTS[rt].mjcf_model_path
363370
robot_cfg.arm_collision_geoms = []

0 commit comments

Comments
 (0)