lerobot conversion and replayer fixes

python/rcs/__main__.py

@@ -57,6 +57,13 @@ def replay(
             help="Parquet dataset directory to replay.",
         ),
     ],
+    output: Annotated[
+        Path,
+        typer.Argument(
+            exists=False,
+            help="Output dir for the new dataset.",
+        ),
+    ],
     headless: Annotated[bool, typer.Option(help="Whether to run without GUI.")] = True,
     frequency: Annotated[int, typer.Option(help="Simulation frequency to use during replay.")] = 30,
     relative_to: Annotated[
@@ -74,6 +81,7 @@ def replay(
 ):
     replay_dataset(
         dataset=dataset,
+        output=output,
         headless=headless,
         frequency=frequency,
         relative_to=relative_to,
@@ -101,7 +109,7 @@ def lerobot_convert(
         str, typer.Option(help="LeRobot repo id metadata. Example: --repo-id myorg/grasp_v2")
     ] = DEFAULT_REPO_ID,
     robot_type: Annotated[
-        str, typer.Option(help="Robot type for metadata and IK model lookup. Example: --robot-type fr3")
+        str, typer.Option(help="Robot type for metadata and IK model lookup. Example: --robot-type FR3")
     ] = DEFAULT_ROBOT_TYPE,
     fps: Annotated[int, typer.Option(help="Dataset frames per second. Example: --fps 30")] = DEFAULT_FPS,
     robot_keys: Annotated[
@@ -138,6 +146,7 @@ def lerobot_convert(
     ] = DEFAULT_PER_ROBOT_ARM_DIM,
     success: Annotated[bool, typer.Option(help="Only include successful episodes. Example: --success")] = True,
     n: Annotated[int, typer.Option(help="Maximum number of episodes to convert. -1 means all. Example: --n 50")] = -1,
+    video_encoding: Annotated[bool, typer.Option(help="Should the image data be video encoded")] = False,
 ):
     cameras = camera_specs_to_configs(camera_specs) if camera_specs is not None else list(DEFAULT_CAMERAS)
     run_conversion(
@@ -154,6 +163,7 @@ def lerobot_convert(
         per_robot_arm_dim=per_robot_arm_dim,
         success=success,
         n=n,
+        video_encoding=video_encoding,
     )
 
 

python/rcs/envs/storage_wrapper.py

@@ -24,7 +24,8 @@ def __init__(
         self,
         env: gym.Env,
         base_dir: str,
-        instruction: str,
+        instruction: str | None = None,
+        allow_wrapper_instruction: bool = True,
         batch_size: int = 32,
         schema: Optional[pa.Schema] = None,
         always_record: bool = False,
@@ -93,6 +94,7 @@ def __init__(
         self._prev_action = None
         self._prev_absolute_action = None
         self.success_from_env = success_from_env
+        self.allow_wrapper_instruction = allow_wrapper_instruction
 
         self.thread_pool = ThreadPoolExecutor()
         self.queue: Queue[pa.Table | pa.RecordBatch] = Queue(maxsize=2)
@@ -259,6 +261,8 @@ def step(self, action):
             self._flatten_arrays(obs)
             if info.get("success") and self.success_from_env:
                 self.success()
+            if info.get("instruction") is not None and self.allow_wrapper_instruction:
+                self.instruction = info.get("instruction")
 
             frame = {
                 "obs": obs,

python/rcs/lerobot_joint_converter.py

@@ -21,7 +21,7 @@
 ]
 DEFAULT_HF_DATA_DIR = "data_lerobot_joint_simple"
 DEFAULT_REPO_ID = "rcs/grasp_joint_simple"
-DEFAULT_ROBOT_TYPE = "fr3"
+DEFAULT_ROBOT_TYPE = "FR3"
 DEFAULT_FPS = 30
 DEFAULT_ROBOT_KEYS = ["left", "right"]
 DEFAULT_JOINTS = False
@@ -97,6 +97,7 @@ def __init__(
         cameras: list[CamConversionConfig] | None = None,
         image_batch_size: int = DEFAULT_IMAGE_BATCH_SIZE,
         per_robot_arm_dim: int = DEFAULT_PER_ROBOT_ARM_DIM,
+        video_encoding: bool = False,
     ):
         self.root = Path(root)
         self.conn = duckdb.connect()
@@ -113,6 +114,7 @@ def __init__(
         self.per_robot_state_dim = self.per_robot_arm_dim + 1
         self.state_dim = len(self.robot_keys) * self.per_robot_state_dim
         self.source_sql = self._build_source_sql(self.dataset_paths)
+        self.video_encoding = video_encoding
 
         self.tcp_offset = rcs.GRIPPER_OFFSETS[self.gripper_type]
         self.ik = rcs.common.Pin(
@@ -126,10 +128,10 @@ def __init__(
             robot_type=self.robot_type.id,
             root=self.root,
             fps=self.fps,
-            use_videos=False,
+            use_videos=self.video_encoding,
             features=self._build_features(),
-            image_writer_threads=0,
-            image_writer_processes=0,
+            image_writer_threads=10,
+            image_writer_processes=5,
         )
 
     def _build_features(self) -> dict[str, dict[str, Any]]:
@@ -140,7 +142,7 @@ def _build_features(self) -> dict[str, dict[str, Any]]:
 
         features = {
             camera.dataset_key: {
-                "dtype": "image",
+                "dtype": "video" if self.video_encoding else "image",
                 "shape": (*camera.resolution, 3),
                 "names": ["height", "width", "channel"],
             }
@@ -409,6 +411,7 @@ def run_conversion(
     per_robot_arm_dim: int = DEFAULT_PER_ROBOT_ARM_DIM,
     success: bool = True,
     n: int = -1,
+    video_encoding: bool = False,
 ) -> None:
     robot_type_converted = RobotType(robot_type)
     gripper_type_converted = GripperType(gripper_type)
@@ -424,6 +427,7 @@ def run_conversion(
         cameras=cameras,
         image_batch_size=image_batch_size,
         per_robot_arm_dim=per_robot_arm_dim,
+        video_encoding=video_encoding,
     )
     converter.generate_examples(success=success, n=n)
 

python/rcs/sim/replayer.py

@@ -13,8 +13,6 @@
 from rcs.envs.scenes import SimEnvCreator
 from rcs.envs.storage_wrapper import StorageWrapper
 
-DATASET_PATH = "recorded_iris"
-
 
 @dataclass(frozen=True)
 class RecordedSimStep:
@@ -118,6 +116,7 @@ def replay_trajectory(env: gym.Env, recorded_steps: list[RecordedSimStep], headl
 
 def replay(
     dataset: Path | str,
+    output: Path | str,
     headless: bool = True,
     frequency: int = 30,
     relative_to: str = RelativeTo.CONFIGURED_ORIGIN.name,
@@ -148,12 +147,13 @@ def replay(
     env_rel = sc.create_env(sim_cfg_data)
     env_rel = StorageWrapper(
         env_rel,
-        DATASET_PATH,
+        str(output),
         "",
         batch_size=32,
         max_rows_per_group=100,
         max_rows_per_file=1000,
         always_record=True,
+        allow_wrapper_instruction=False,
     )
     try:
         for uuid in uuids: