From cb9be661359beb85afffd1af8cdb7354ad6779fe Mon Sep 17 00:00:00 2001 From: aditya-neuraco Date: Fri, 12 Jun 2026 10:31:39 +0100 Subject: [PATCH 1/3] feat: add nyu rot dataset --- neuracore-dictionary.txt | 5 +- .../importer/config/nyu_rot_dataset.yaml | 98 ++++++ neuracore/importer/config/nyu_rot_rlds.yaml | 61 ++++ scripts/annotate_gripper.py | 277 +++++++++++++++ scripts/gripper_states.json | 16 + scripts/import_nyu_rot.py | 314 ++++++++++++++++++ 6 files changed, 770 insertions(+), 1 deletion(-) create mode 100644 neuracore/importer/config/nyu_rot_dataset.yaml create mode 100644 neuracore/importer/config/nyu_rot_rlds.yaml create mode 100644 scripts/annotate_gripper.py create mode 100644 scripts/gripper_states.json create mode 100644 scripts/import_nyu_rot.py diff --git a/neuracore-dictionary.txt b/neuracore-dictionary.txt index bd433f9e3..d9ea40742 100644 --- a/neuracore-dictionary.txt +++ b/neuracore-dictionary.txt @@ -399,4 +399,7 @@ octomap delenv hookwrapper makereport -tryfirst \ No newline at end of file +tryfirst +haldar +osfstorage +robotgym diff --git a/neuracore/importer/config/nyu_rot_dataset.yaml b/neuracore/importer/config/nyu_rot_dataset.yaml new file mode 100644 index 000000000..ac8344707 --- /dev/null +++ b/neuracore/importer/config/nyu_rot_dataset.yaml @@ -0,0 +1,98 @@ +input_dataset_name: lerobot/nyu_rot_dataset +dataset_type: LEROBOT +output_dataset: + name: NYU ROT + tags: [xarm, manipulation, pick-and-place, rot, nyu, imitation-learning] + description: "NYU ROT (Regularized Optimal Transport) dataset collected on a Ufactory xArm 7 robot with xArm Gripper. Contains 14 episodes of pick-and-place manipulation tasks with RGB observations and proprioceptive motor state, used in the paper 'Watch and Match: Supercharging Imitation with Regularized Optimal Transport' (Haldar et al., 2023)." + +robot: + name: nyu_rot_xarm7 + urdf_path: "" + +frequency: 5 + +data_import_config: + RGB_IMAGES: + source: observation.images + format: + image_convention: CHANNELS_FIRST + order_of_channels: RGB + normalized_pixel_values: TRUE + mapping: + - name: image + source_name: image + + JOINT_POSITIONS: + source: observation.state + format: + angle_units: RADIANS + mapping: + - name: joint1 + index: 0 + - name: joint2 + index: 1 + - name: joint3 + index: 2 + - name: joint4 + index: 3 + - name: joint5 + index: 4 + - name: joint6 + index: 5 + - name: joint7 + index: 6 + + JOINT_TARGET_POSITIONS: + source: action + format: + action_type: ABSOLUTE + angle_units: RADIANS + mapping: + - name: joint1 + index: 0 + - name: joint2 + index: 1 + - name: joint3 + index: 2 + - name: joint4 + index: 3 + - name: joint5 + index: 4 + - name: joint6 + index: 5 + - name: joint7 + index: 6 + + END_EFFECTOR_POSES: + source: observation.state + format: + ee_pose_input_type: JOINT_POSITIONS + mapping: + - name: link_eef + + VISUAL_JOINT_POSITIONS: + source: observation.state + format: + visual_joint_input_type: CUSTOM + mapping: + - name: joint1 + index: 0 + - name: joint2 + index: 1 + - name: joint3 + index: 2 + - name: joint4 + index: 3 + - name: joint5 + index: 4 + - name: joint6 + index: 5 + - name: joint7 + index: 6 + + LANGUAGE: + source: language_instruction + format: + language_type: STRING + mapping: + - name: instruction diff --git a/neuracore/importer/config/nyu_rot_rlds.yaml b/neuracore/importer/config/nyu_rot_rlds.yaml new file mode 100644 index 000000000..9f64a7ea7 --- /dev/null +++ b/neuracore/importer/config/nyu_rot_rlds.yaml @@ -0,0 +1,61 @@ +input_dataset_name: nyu_rot_dataset_converted_externally_to_rlds +dataset_type: RLDS + +output_dataset: + name: nyu_rot + tags: [xarm, table, manipulation, imitation-learning] + description: "NYU ROT dataset from the Watch and Match (Regularized Optimal Transport) project. Real-world tabletop manipulation demonstrations collected on a UFactory xArm, with a single 84x84 RGB camera view, end-effector state and delta end-effector actions across 14 manipulation tasks." + +robot: + name: xarm + urdf_path: "" + overwrite_existing: true + +frequency: 3.0 + +data_import_config: + RGB_IMAGES: + source: observation + format: + image_convention: CHANNELS_LAST + order_of_channels: RGB + mapping: + - name: image + source_name: image + + END_EFFECTOR_POSES: + source: observation.state + format: + pose_type: POSITION_ORIENTATION + orientation: + type: EULER + euler_order: XYZ + angle_units: RADIANS + mapping: + - name: end_effector_pose + index_range: + start: 0 + end: 6 + + PARALLEL_GRIPPER_OPEN_AMOUNTS: + source: observation.state + format: + invert_gripper_amount: true # dataset: 0-open, 1-closed + mapping: + - name: gripper + index: 6 + + PARALLEL_GRIPPER_TARGET_OPEN_AMOUNTS: + source: action + format: + invert_gripper_amount: true # dataset: 0-open, 1-closed + mapping: + - name: gripper + index: 6 + + LANGUAGE: + format: + language_type: BYTES + mapping: + - name: instruction + source_name: language_instruction diff --git a/scripts/annotate_gripper.py b/scripts/annotate_gripper.py new file mode 100644 index 000000000..e34b6cfdd --- /dev/null +++ b/scripts/annotate_gripper.py @@ -0,0 +1,277 @@ +"""GUI tool to manually annotate gripper state (open/closed) for each NYU ROT episode. + +Loads pkl files, shows image frames, lets user toggle open/closed per episode, +saves result to gripper_states.json. +""" + +import json +import pickle +import sys +import zipfile +from pathlib import Path + +import numpy as np +from PyQt6.QtCore import Qt, QTimer +from PyQt6.QtGui import QImage, QPixmap +from PyQt6.QtWidgets import ( + QApplication, + QButtonGroup, + QHBoxLayout, + QLabel, + QMainWindow, + QPushButton, + QRadioButton, + QSlider, + QVBoxLayout, + QWidget, +) + +EXPERT_DEMOS_ZIP = Path.home() / "Downloads" / "osfstorage-archive.zip" +EXTRACT_DIR = Path("/tmp/nyu_rot_extract") +OUTPUT_FILE = Path(__file__).parent / "gripper_states.json" + +TASK_KEYS = [ + "RobotEraseBoard-v1/expert_demos", + "RobotHangHanger-v1/expert_demos", + "RobotReach-v1/expert_demos", + "RobotDoorClose-v1/expert_demos", + "RobotCupStacking-v1/expert_demos", + "RobotTurnKnob-v1/expert_demos", + "RobotInsertPeg-v1/expert_demos_easy", + "RobotInsertPeg-v1/expert_demos_medium", + "RobotInsertPeg-v1/expert_demos_hard", + "RobotButtonPress-v1/expert_demos", + "RobotHangBag-v1/expert_demos", + "RobotBoxOpen-v1/expert_demos", + "RobotPour-v1/expert_demos", + "RobotHangMug-v1/expert_demos", +] + +# Defaults from import script +DEFAULTS = { + "RobotEraseBoard-v1/expert_demos": True, + "RobotHangHanger-v1/expert_demos": True, + "RobotReach-v1/expert_demos": False, + "RobotDoorClose-v1/expert_demos": True, + "RobotCupStacking-v1/expert_demos": True, + "RobotTurnKnob-v1/expert_demos": True, + "RobotInsertPeg-v1/expert_demos_easy": True, + "RobotInsertPeg-v1/expert_demos_medium": True, + "RobotInsertPeg-v1/expert_demos_hard": True, + "RobotButtonPress-v1/expert_demos": True, + "RobotHangBag-v1/expert_demos": True, + "RobotBoxOpen-v1/expert_demos": False, + "RobotPour-v1/expert_demos": True, + "RobotHangMug-v1/expert_demos": True, +} + + +def load_episodes() -> dict[str, np.ndarray]: + """Load images for all episodes from pkl files.""" + EXTRACT_DIR.mkdir(parents=True, exist_ok=True) + demos_zip = EXTRACT_DIR / "expert_demos.zip" + if not demos_zip.exists(): + print("Extracting expert_demos.zip...") + with zipfile.ZipFile(EXPERT_DEMOS_ZIP) as zf: + zf.extract("expert_demos.zip", EXTRACT_DIR) + robotgym_dir = EXTRACT_DIR / "expert_demos" / "robotgym" + if not robotgym_dir.exists(): + print("Extracting robotgym pkl files...") + with zipfile.ZipFile(demos_zip) as zf: + for name in zf.namelist(): + if name.startswith("expert_demos/robotgym/"): + zf.extract(name, EXTRACT_DIR) + + episodes = {} + for task_dir in sorted(robotgym_dir.iterdir()): + for pkl_path in sorted(task_dir.iterdir()): + key = f"{task_dir.name}/{pkl_path.stem}" + if key not in TASK_KEYS: + continue + with open(pkl_path, "rb") as f: + images, *_ = pickle.load(f) + # images: (n_eps, n_steps, 3, H, W) — take episode 0, CHW→HWC + episodes[key] = np.transpose(images[0], (0, 2, 3, 1)) # (steps, H, W, 3) + return episodes + + +def ndarray_to_pixmap(frame: np.ndarray, scale: int = 4) -> QPixmap: + h, w, c = frame.shape + img = QImage(frame.tobytes(), w, h, w * c, QImage.Format.Format_RGB888) + return QPixmap.fromImage(img).scaled( + w * scale, h * scale, Qt.AspectRatioMode.KeepAspectRatio + ) + + +class GripperAnnotator(QMainWindow): + def __init__(self, episodes: dict[str, np.ndarray], saved: dict[str, bool]) -> None: + super().__init__() + self.episodes = episodes + self.keys = TASK_KEYS + self.states: dict[str, bool] = { + k: saved.get(k, DEFAULTS.get(k, False)) for k in self.keys + } + self.current = 0 + self.playing = False + self.frame_idx = 0 + self._timer = QTimer() + self._timer.timeout.connect(self._next_frame) + + self.setWindowTitle("NYU ROT Gripper Annotator") + self._build_ui() + self._load_episode(0) + + def _build_ui(self) -> None: + central = QWidget() + self.setCentralWidget(central) + root = QVBoxLayout(central) + + # Episode label + self.ep_label = QLabel() + self.ep_label.setAlignment(Qt.AlignmentFlag.AlignCenter) + self.ep_label.setStyleSheet("font-size: 14px; font-weight: bold;") + root.addWidget(self.ep_label) + + # Progress + self.progress_label = QLabel() + self.progress_label.setAlignment(Qt.AlignmentFlag.AlignCenter) + root.addWidget(self.progress_label) + + # Image display + self.image_label = QLabel() + self.image_label.setAlignment(Qt.AlignmentFlag.AlignCenter) + root.addWidget(self.image_label) + + # Frame slider + self.slider = QSlider(Qt.Orientation.Horizontal) + self.slider.valueChanged.connect(self._on_slider) + root.addWidget(self.slider) + + self.frame_label = QLabel() + self.frame_label.setAlignment(Qt.AlignmentFlag.AlignCenter) + root.addWidget(self.frame_label) + + # Play/pause + play_row = QHBoxLayout() + self.play_btn = QPushButton("▶ Play") + self.play_btn.clicked.connect(self._toggle_play) + play_row.addWidget(self.play_btn) + root.addLayout(play_row) + + # Gripper toggle + grip_row = QHBoxLayout() + self.btn_group = QButtonGroup() + self.radio_closed = QRadioButton("Closed (gripper holding object)") + self.radio_open = QRadioButton("Open (gripper not holding)") + self.btn_group.addButton(self.radio_closed) + self.btn_group.addButton(self.radio_open) + self.radio_closed.toggled.connect(self._on_gripper_changed) + grip_row.addWidget(self.radio_closed) + grip_row.addWidget(self.radio_open) + root.addLayout(grip_row) + + # Nav buttons + nav_row = QHBoxLayout() + self.prev_btn = QPushButton("← Previous") + self.prev_btn.clicked.connect(self._prev) + self.next_btn = QPushButton("Next →") + self.next_btn.clicked.connect(self._next) + self.save_btn = QPushButton("💾 Save") + self.save_btn.setStyleSheet( + "background-color: #4CAF50; color: white; font-weight: bold;" + ) + self.save_btn.clicked.connect(self._save) + nav_row.addWidget(self.prev_btn) + nav_row.addWidget(self.save_btn) + nav_row.addWidget(self.next_btn) + root.addLayout(nav_row) + + # Saved indicator + self.saved_label = QLabel() + self.saved_label.setAlignment(Qt.AlignmentFlag.AlignCenter) + root.addWidget(self.saved_label) + + def _load_episode(self, idx: int) -> None: + self._timer.stop() + self.playing = False + self.play_btn.setText("▶ Play") + self.current = idx + key = self.keys[idx] + self.frames = self.episodes[key] + self.frame_idx = 0 + n = len(self.frames) + + self.ep_label.setText(key) + self.progress_label.setText(f"Episode {idx + 1} / {len(self.keys)}") + self.slider.setMaximum(n - 1) + self.slider.setValue(0) + + closed = self.states[key] + self.radio_closed.blockSignals(True) + self.radio_open.blockSignals(True) + (self.radio_closed if closed else self.radio_open).setChecked(True) + self.radio_closed.blockSignals(False) + self.radio_open.blockSignals(False) + + self.prev_btn.setEnabled(idx > 0) + self.next_btn.setEnabled(idx < len(self.keys) - 1) + self._show_frame(0) + + def _show_frame(self, idx: int) -> None: + self.frame_idx = idx + self.image_label.setPixmap(ndarray_to_pixmap(self.frames[idx])) + self.frame_label.setText(f"Frame {idx + 1} / {len(self.frames)}") + + def _on_slider(self, val: int) -> None: + self._show_frame(val) + + def _toggle_play(self) -> None: + if self.playing: + self._timer.stop() + self.playing = False + self.play_btn.setText("▶ Play") + else: + self.playing = True + self.play_btn.setText("⏸ Pause") + self._timer.start(200) # 5 Hz + + def _next_frame(self) -> None: + next_idx = (self.frame_idx + 1) % len(self.frames) + self.slider.setValue(next_idx) + + def _on_gripper_changed(self) -> None: + self.states[self.keys[self.current]] = self.radio_closed.isChecked() + + def _prev(self) -> None: + if self.current > 0: + self._load_episode(self.current - 1) + + def _next(self) -> None: + if self.current < len(self.keys) - 1: + self._load_episode(self.current + 1) + + def _save(self) -> None: + with open(OUTPUT_FILE, "w") as f: + json.dump(self.states, f, indent=2) + self.saved_label.setText(f"Saved to {OUTPUT_FILE}") + self.saved_label.setStyleSheet("color: green;") + + +def main() -> None: + print("Loading episodes...") + episodes = load_episodes() + + saved = {} + if OUTPUT_FILE.exists(): + with open(OUTPUT_FILE) as f: + saved = json.load(f) + print(f"Loaded existing annotations from {OUTPUT_FILE}") + + app = QApplication(sys.argv) + window = GripperAnnotator(episodes, saved) + window.show() + sys.exit(app.exec()) + + +if __name__ == "__main__": + main() diff --git a/scripts/gripper_states.json b/scripts/gripper_states.json new file mode 100644 index 000000000..126504c5a --- /dev/null +++ b/scripts/gripper_states.json @@ -0,0 +1,16 @@ +{ + "RobotEraseBoard-v1/expert_demos": true, + "RobotHangHanger-v1/expert_demos": true, + "RobotReach-v1/expert_demos": true, + "RobotDoorClose-v1/expert_demos": true, + "RobotCupStacking-v1/expert_demos": true, + "RobotTurnKnob-v1/expert_demos": true, + "RobotInsertPeg-v1/expert_demos_easy": true, + "RobotInsertPeg-v1/expert_demos_medium": true, + "RobotInsertPeg-v1/expert_demos_hard": true, + "RobotButtonPress-v1/expert_demos": true, + "RobotHangBag-v1/expert_demos": true, + "RobotBoxOpen-v1/expert_demos": false, + "RobotPour-v1/expert_demos": true, + "RobotHangMug-v1/expert_demos": true +} \ No newline at end of file diff --git a/scripts/import_nyu_rot.py b/scripts/import_nyu_rot.py new file mode 100644 index 000000000..9dba9901a --- /dev/null +++ b/scripts/import_nyu_rot.py @@ -0,0 +1,314 @@ +"""Import NYU ROT dataset directly into Neuracore from original pkl files. + +Data format (from transform_pickle.py + demo_class.py): + [0] images: (n_eps, n_steps, 3, 84, 84) uint8 + [1] states: (n_eps, n_steps, 3) float32 — EE pos (x,y,z) in decimeters + (n_eps, n_steps, 4) for Pour — EE pos + joint7 servo angle (degrees) + [2] actions: (n_eps, n_steps, 3) float32 — delta EE pos / 0.25, range [-1,1] + (n_eps, n_steps, 4) for Pour — delta EE pos / 0.25 + delta joint7 / 40 + [3] rewards: (n_eps, n_steps) float32 — 0 during task, 1 at success/end + +EE orientation is FIXED per task class (pitch/roll set at robot init, never changes). +xArm units: positions in decimeters (0.1 m). Angles in degrees. +Pinocchio expects SI units (meters, radians). +""" + +import json +import os +import pickle +import subprocess +import sys +import time +import zipfile +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +from scipy.spatial.transform import Rotation as R + +import neuracore as nc +from neuracore.importer.core.robot_utils import RobotUtils + +EXPERT_DEMOS_ZIP = Path.home() / "Downloads" / "osfstorage-archive.zip" +EXTRACT_DIR = Path("/tmp/nyu_rot_extract") +URDF_PATH = Path.home() / "xarm_ros2" / "xarm7.urdf" +FPS = 5 +ACTION_SCALE_POS = 0.25 # stored_action * this = delta EE pos in decimeters +ACTION_SCALE_ROT = 40.0 # stored_action[3] * this = delta joint7 in degrees +EE_FRAME = "link_eef" +GRIPPER_NAME = "gripper" +GRIPPER_OPEN_RAD = 0.0 # drive_joint lower limit = open +GRIPPER_CLOSED_RAD = 0.85 # drive_joint upper limit = closed +GRIPPER_JOINT_NAMES = { + "drive_joint", + "left_finger_joint", + "left_inner_knuckle_joint", + "right_outer_knuckle_joint", + "right_finger_joint", + "right_inner_knuckle_joint", +} + + +@dataclass +class TaskMeta: + instruction: str + roll: float # degrees, xArm extrinsic RPY + pitch: float + yaw: float + gripper_closed: bool # constant throughout episode + has_wrist: bool = False # Pour: joint7 explicitly recorded + + +# Orientation from demo_class.py defaults per class: +# Reach/BaseClass default: pitch=0, roll=180 +# HorizontalReach: pitch=-90, roll=180 +# SideReach: pitch=0, roll=90 +# Pour: pitch=90, roll=-90 +TASK_META: dict[str, TaskMeta] = { + "RobotEraseBoard-v1/expert_demos": TaskMeta( + "erase the board", 180, 0, 0, gripper_closed=True + ), + "RobotHangHanger-v1/expert_demos": TaskMeta( + "hang the hanger on the rod", 180, -90, 0, gripper_closed=True + ), + "RobotReach-v1/expert_demos": TaskMeta( + "reach the blue mark on the table", 180, 0, 0, gripper_closed=False + ), + "RobotDoorClose-v1/expert_demos": TaskMeta( + "close the door", 180, 0, 0, gripper_closed=True + ), + "RobotCupStacking-v1/expert_demos": TaskMeta( + "stack the cups", 180, 0, 0, gripper_closed=True + ), + "RobotTurnKnob-v1/expert_demos": TaskMeta( + "turn the knob", 180, 0, 0, gripper_closed=True + ), + "RobotInsertPeg-v1/expert_demos_easy": TaskMeta( + "insert the peg in the cup (easy)", 180, 0, 0, gripper_closed=True + ), + "RobotInsertPeg-v1/expert_demos_medium": TaskMeta( + "insert the peg in the cup (medium)", 180, 0, 0, gripper_closed=True + ), + "RobotInsertPeg-v1/expert_demos_hard": TaskMeta( + "insert the peg in the cup (hard)", 180, 0, 0, gripper_closed=True + ), + "RobotButtonPress-v1/expert_demos": TaskMeta( + "press the button", 180, 0, 0, gripper_closed=True + ), + "RobotHangBag-v1/expert_demos": TaskMeta( + "hang the bag on the hook", 180, -90, 0, gripper_closed=True + ), + "RobotBoxOpen-v1/expert_demos": TaskMeta( + "open the box", 90, 0, 0, gripper_closed=False + ), + "RobotPour-v1/expert_demos": TaskMeta( + "pour the almonds into the cup", -90, 90, 0, gripper_closed=True, has_wrist=True + ), + "RobotHangMug-v1/expert_demos": TaskMeta( + "hang the mug on the hook", 180, -90, 0, gripper_closed=True + ), +} + + +def rpy_to_quat(roll_deg: float, pitch_deg: float, yaw_deg: float) -> np.ndarray: + """xArm extrinsic RPY (degrees) → quaternion [qx, qy, qz, qw].""" + return R.from_euler("xyz", [roll_deg, pitch_deg, yaw_deg], degrees=True).as_quat() + + +def build_ee_pose(pos_dm: np.ndarray, quat_xyzw: np.ndarray) -> np.ndarray: + """Combine EE position (decimeters) + quaternion into 7-element pose array.""" + return np.array([*pos_dm, *quat_xyzw], dtype=np.float64) + + +def extract_data() -> Path: + EXTRACT_DIR.mkdir(parents=True, exist_ok=True) + demos_zip = EXTRACT_DIR / "expert_demos.zip" + if not demos_zip.exists(): + print("Extracting expert_demos.zip from archive...") + with zipfile.ZipFile(EXPERT_DEMOS_ZIP) as zf: + zf.extract("expert_demos.zip", EXTRACT_DIR) + robotgym_dir = EXTRACT_DIR / "expert_demos" / "robotgym" + if not robotgym_dir.exists(): + print("Extracting robotgym pkl files...") + with zipfile.ZipFile(demos_zip) as zf: + for name in zf.namelist(): + if name.startswith("expert_demos/robotgym/"): + zf.extract(name, EXTRACT_DIR) + return robotgym_dir + + +def collect_episodes(robotgym_dir: Path): + """Yield (key, meta, images, states, actions, rewards) per episode.""" + for task_dir in sorted(robotgym_dir.iterdir()): + for pkl_path in sorted(task_dir.iterdir()): + key = f"{task_dir.name}/{pkl_path.stem}" + meta = TASK_META.get(key) + if meta is None: + print(f" Skipping unknown key: {key}") + continue + with open(pkl_path, "rb") as f: + images, states, actions, rewards = pickle.load(f) + for ep_idx in range(images.shape[0]): + yield key, meta, images[ep_idx], states[ep_idx], actions[ + ep_idx + ], rewards[ep_idx] + + +def _restart_daemon() -> None: + daemon_cmd = [sys.executable, "-m", "neuracore.data_daemon"] + subprocess.run([*daemon_cmd, "stop"], check=False) + time.sleep(2) + subprocess.run([*daemon_cmd, "launch", "--background"], check=True) + time.sleep(2) + print("Daemon restarted.") + + +def main(): + robotgym_dir = extract_data() + robot_utils = RobotUtils(str(URDF_PATH), os.path.dirname(str(URDF_PATH))) + + _restart_daemon() + nc.login() + nc.connect_robot( + robot_name="Ufactory xArm 7", + urdf_path=str(URDF_PATH), + overwrite=True, + shared=True, + ) + + from neuracore.core.data.dataset import Dataset + + existing = Dataset.get_by_name("NYU ROT", non_exist_ok=True) + if existing is not None: + print(f"Deleting existing dataset '{existing.id}' before recreating...") + existing.delete() + + nc.create_dataset( + name="NYU ROT", + tags=[ + "xarm", + "manipulation", + "pick-and-place", + "rot", + "nyu", + "imitation-learning", + ], + description=( + "NYU ROT (Regularized Optimal Transport) dataset collected on a " + "Ufactory xArm 7 robot. 14 episodes across 12 manipulation tasks. " + "From 'Watch and Match: Supercharging Imitation with Regularized " + "Optimal Transport' (Haldar et al., 2023)." + ), + shared=True, + ) + + episodes = list(collect_episodes(robotgym_dir)) + print(f"Found {len(episodes)} episodes") + + gripper_annotations_path = Path(__file__).parent / "gripper_states.json" + gripper_annotations: dict[str, bool] = {} + if gripper_annotations_path.exists(): + with open(gripper_annotations_path) as f: + gripper_annotations = json.load(f) + print(f"Using gripper annotations from {gripper_annotations_path}") + + gripper_open = 1.0 + gripper_closed = 0.0 + + for ep_idx, (key, meta, images, states, actions, rewards) in enumerate(episodes): + n_steps = images.shape[0] + print( + f"Episode {ep_idx + 1}/{len(episodes)}: {meta.instruction} " + f"({n_steps} steps)" + ) + + quat = rpy_to_quat(meta.roll, meta.pitch, meta.yaw) + is_closed = gripper_annotations.get(key, meta.gripper_closed) + gripper_val = gripper_closed if is_closed else gripper_open + gripper_rad = GRIPPER_CLOSED_RAD if is_closed else GRIPPER_OPEN_RAD + gripper_joints = { + "drive_joint": gripper_rad, + "left_finger_joint": gripper_rad, + "left_inner_knuckle_joint": gripper_rad, + "right_outer_knuckle_joint": gripper_rad, + "right_finger_joint": gripper_rad, + "right_inner_knuckle_joint": gripper_rad, + } + + nc.start_recording() + t = time.time() + dt = 1.0 / FPS + prev_ik = None + prev_target_ik = None + + for step in range(n_steps): + pos_m = states[step, :3] / 10.0 # decimeters → meters + + # IK for current EE pose + try: + joint_pos = robot_utils.end_effector_to_joint_positions( + build_ee_pose(pos_m, quat), EE_FRAME, prev_ik + ) + prev_ik = list(joint_pos.values()) + arm_joints = { + k: v for k, v in joint_pos.items() if k not in GRIPPER_JOINT_NAMES + } + nc.log_joint_positions(arm_joints, timestamp=t) + nc.log_visual_joint_positions(arm_joints, timestamp=t) + except ValueError: + print(f" IK failed at step {step}, skipping joint positions") + + # IK for target EE pose + target_pos_m = pos_m + actions[step, :3] * ACTION_SCALE_POS / 10.0 + try: + target_joint_pos = robot_utils.end_effector_to_joint_positions( + build_ee_pose(target_pos_m, quat), EE_FRAME, prev_target_ik + ) + prev_target_ik = list(target_joint_pos.values()) + arm_target_joints = { + k: v + for k, v in target_joint_pos.items() + if k not in GRIPPER_JOINT_NAMES + } + nc.log_joint_target_positions(arm_target_joints, timestamp=t) + except ValueError: + print(f" Target IK failed at step {step}, skipping") + + # EE pose (meters) + nc.log_end_effector_pose(EE_FRAME, build_ee_pose(pos_m, quat), timestamp=t) + + # Target EE pose + nc.log_pose("ee_target", build_ee_pose(target_pos_m, quat), timestamp=t) + + # Gripper (binary constant per episode) + nc.log_parallel_gripper_open_amount(GRIPPER_NAME, gripper_val, timestamp=t) + nc.log_parallel_gripper_target_open_amount( + GRIPPER_NAME, gripper_val, timestamp=t + ) + nc.log_joint_positions(gripper_joints, timestamp=t) + nc.log_visual_joint_positions(gripper_joints, timestamp=t) + nc.log_joint_target_positions(gripper_joints, timestamp=t) + + # Pour: wrist servo angle (joint7) in radians — overrides IK joint7 + if meta.has_wrist: + j7_rad = float(np.deg2rad(states[step, 3])) + target_j7 = j7_rad + float( + np.deg2rad(actions[step, 3] * ACTION_SCALE_ROT) + ) + nc.log_joint_positions({"joint7": j7_rad}, timestamp=t) + nc.log_visual_joint_positions({"joint7": j7_rad}, timestamp=t) + nc.log_joint_target_positions({"joint7": target_j7}, timestamp=t) + + nc.log_rgb("image", np.transpose(images[step], (1, 2, 0)), timestamp=t) + nc.log_language("instruction", meta.instruction, timestamp=t) + + t += dt + + nc.stop_recording() + print(" Done") + + print("All episodes imported.") + + +if __name__ == "__main__": + main() From 63f95f616e4e5bbc5597d5a38166c58a3ec302ca Mon Sep 17 00:00:00 2001 From: ypang-neuraco Date: Mon, 15 Jun 2026 11:58:49 +0100 Subject: [PATCH 2/3] fix: ik smoothness --- scripts/annotate_gripper.py | 29 ++++++++++++------- scripts/import_nyu_rot.py | 58 +++++++++++++++++++++++++------------ 2 files changed, 58 insertions(+), 29 deletions(-) diff --git a/scripts/annotate_gripper.py b/scripts/annotate_gripper.py index e34b6cfdd..bc5f289d7 100644 --- a/scripts/annotate_gripper.py +++ b/scripts/annotate_gripper.py @@ -69,18 +69,27 @@ def load_episodes() -> dict[str, np.ndarray]: """Load images for all episodes from pkl files.""" EXTRACT_DIR.mkdir(parents=True, exist_ok=True) - demos_zip = EXTRACT_DIR / "expert_demos.zip" - if not demos_zip.exists(): - print("Extracting expert_demos.zip...") - with zipfile.ZipFile(EXPERT_DEMOS_ZIP) as zf: - zf.extract("expert_demos.zip", EXTRACT_DIR) robotgym_dir = EXTRACT_DIR / "expert_demos" / "robotgym" if not robotgym_dir.exists(): - print("Extracting robotgym pkl files...") - with zipfile.ZipFile(demos_zip) as zf: - for name in zf.namelist(): - if name.startswith("expert_demos/robotgym/"): - zf.extract(name, EXTRACT_DIR) + with zipfile.ZipFile(EXPERT_DEMOS_ZIP) as zf: + names = zf.namelist() + if "expert_demos.zip" in names: + print("Extracting expert_demos.zip...") + zf.extract("expert_demos.zip", EXTRACT_DIR) + with zipfile.ZipFile(EXTRACT_DIR / "expert_demos.zip") as inner: + print("Extracting robotgym pkl files...") + for name in inner.namelist(): + if name.startswith("expert_demos/robotgym/"): + inner.extract(name, EXTRACT_DIR) + elif any(name.startswith("expert_demos/robotgym/") for name in names): + print("Extracting robotgym pkl files...") + for name in names: + if name.startswith("expert_demos/robotgym/"): + zf.extract(name, EXTRACT_DIR) + else: + raise FileNotFoundError( + f"Could not find expert_demos/robotgym/ in {EXPERT_DEMOS_ZIP}" + ) episodes = {} for task_dir in sorted(robotgym_dir.iterdir()): diff --git a/scripts/import_nyu_rot.py b/scripts/import_nyu_rot.py index 9dba9901a..6441eff12 100644 --- a/scripts/import_nyu_rot.py +++ b/scripts/import_nyu_rot.py @@ -29,9 +29,9 @@ import neuracore as nc from neuracore.importer.core.robot_utils import RobotUtils -EXPERT_DEMOS_ZIP = Path.home() / "Downloads" / "osfstorage-archive.zip" +EXPERT_DEMOS_ZIP = Path.home() / "Downloads" / "expert_demos.zip" EXTRACT_DIR = Path("/tmp/nyu_rot_extract") -URDF_PATH = Path.home() / "xarm_ros2" / "xarm7.urdf" +URDF_PATH = Path.home() / "Datasets/URDF/xarm/xarm7_with_gripper.urdf" FPS = 5 ACTION_SCALE_POS = 0.25 # stored_action * this = delta EE pos in decimeters ACTION_SCALE_ROT = 40.0 # stored_action[3] * this = delta joint7 in degrees @@ -122,18 +122,29 @@ def build_ee_pose(pos_dm: np.ndarray, quat_xyzw: np.ndarray) -> np.ndarray: def extract_data() -> Path: EXTRACT_DIR.mkdir(parents=True, exist_ok=True) - demos_zip = EXTRACT_DIR / "expert_demos.zip" - if not demos_zip.exists(): - print("Extracting expert_demos.zip from archive...") - with zipfile.ZipFile(EXPERT_DEMOS_ZIP) as zf: - zf.extract("expert_demos.zip", EXTRACT_DIR) robotgym_dir = EXTRACT_DIR / "expert_demos" / "robotgym" - if not robotgym_dir.exists(): - print("Extracting robotgym pkl files...") - with zipfile.ZipFile(demos_zip) as zf: - for name in zf.namelist(): + if robotgym_dir.exists(): + return robotgym_dir + + with zipfile.ZipFile(EXPERT_DEMOS_ZIP) as zf: + names = zf.namelist() + if "expert_demos.zip" in names: + print("Extracting expert_demos.zip from archive...") + zf.extract("expert_demos.zip", EXTRACT_DIR) + with zipfile.ZipFile(EXTRACT_DIR / "expert_demos.zip") as inner: + print("Extracting robotgym pkl files...") + for name in inner.namelist(): + if name.startswith("expert_demos/robotgym/"): + inner.extract(name, EXTRACT_DIR) + elif any(name.startswith("expert_demos/robotgym/") for name in names): + print("Extracting robotgym pkl files...") + for name in names: if name.startswith("expert_demos/robotgym/"): zf.extract(name, EXTRACT_DIR) + else: + raise FileNotFoundError( + f"Could not find expert_demos/robotgym/ in {EXPERT_DEMOS_ZIP}" + ) return robotgym_dir @@ -238,18 +249,21 @@ def main(): nc.start_recording() t = time.time() dt = 1.0 / FPS - prev_ik = None - prev_target_ik = None + prev_ik: dict[str, float] | None = None + prev_target_ik: dict[str, float] | None = None for step in range(n_steps): pos_m = states[step, :3] / 10.0 # decimeters → meters + joint_pos: dict[str, float] | None = None - # IK for current EE pose + # IK for current EE pose. Keep prev_ik as a {joint_name: angle} dict: + # a bare list of arm joint values is not a valid q0 (model.nq includes + # gripper joints), so IK would restart from the limit midpoint each frame. try: joint_pos = robot_utils.end_effector_to_joint_positions( build_ee_pose(pos_m, quat), EE_FRAME, prev_ik ) - prev_ik = list(joint_pos.values()) + prev_ik = joint_pos arm_joints = { k: v for k, v in joint_pos.items() if k not in GRIPPER_JOINT_NAMES } @@ -258,13 +272,15 @@ def main(): except ValueError: print(f" IK failed at step {step}, skipping joint positions") - # IK for target EE pose + # IK for target EE pose. Seed from the current solution when available + # since the target pose is a small delta from the current pose. target_pos_m = pos_m + actions[step, :3] * ACTION_SCALE_POS / 10.0 try: + target_seed = joint_pos if joint_pos is not None else prev_target_ik target_joint_pos = robot_utils.end_effector_to_joint_positions( - build_ee_pose(target_pos_m, quat), EE_FRAME, prev_target_ik + build_ee_pose(target_pos_m, quat), EE_FRAME, target_seed ) - prev_target_ik = list(target_joint_pos.values()) + prev_target_ik = target_joint_pos arm_target_joints = { k: v for k, v in target_joint_pos.items() @@ -298,13 +314,17 @@ def main(): nc.log_joint_positions({"joint7": j7_rad}, timestamp=t) nc.log_visual_joint_positions({"joint7": j7_rad}, timestamp=t) nc.log_joint_target_positions({"joint7": target_j7}, timestamp=t) + if prev_ik is not None: + prev_ik = {**prev_ik, "joint7": j7_rad} + if prev_target_ik is not None: + prev_target_ik = {**prev_target_ik, "joint7": target_j7} nc.log_rgb("image", np.transpose(images[step], (1, 2, 0)), timestamp=t) nc.log_language("instruction", meta.instruction, timestamp=t) t += dt - nc.stop_recording() + nc.stop_recording(wait=True) print(" Done") print("All episodes imported.") From 5cd7ca9deede76920f7ff1bb217274411336f9ea Mon Sep 17 00:00:00 2001 From: ypang-neuraco Date: Mon, 22 Jun 2026 09:54:51 +0100 Subject: [PATCH 3/3] fix: add position offset --- scripts/import_nyu_rot.py | 20 ++++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/scripts/import_nyu_rot.py b/scripts/import_nyu_rot.py index 6441eff12..eab9343eb 100644 --- a/scripts/import_nyu_rot.py +++ b/scripts/import_nyu_rot.py @@ -37,6 +37,17 @@ ACTION_SCALE_ROT = 40.0 # stored_action[3] * this = delta joint7 in degrees EE_FRAME = "link_eef" GRIPPER_NAME = "gripper" +SHARED = True +# Initial joint seed for the first IK solve of each episode (radians). +INITIAL_JOINT_POSITIONS = { + "joint1": 0.0, + "joint2": 0.0, + "joint3": 0.0, + "joint4": 0.0, + "joint5": 0.0, + "joint6": 0.0, + "joint7": 0.0, +} GRIPPER_OPEN_RAD = 0.0 # drive_joint lower limit = open GRIPPER_CLOSED_RAD = 0.85 # drive_joint upper limit = closed GRIPPER_JOINT_NAMES = { @@ -184,7 +195,7 @@ def main(): robot_name="Ufactory xArm 7", urdf_path=str(URDF_PATH), overwrite=True, - shared=True, + shared=SHARED, ) from neuracore.core.data.dataset import Dataset @@ -210,7 +221,7 @@ def main(): "From 'Watch and Match: Supercharging Imitation with Regularized " "Optimal Transport' (Haldar et al., 2023)." ), - shared=True, + shared=SHARED, ) episodes = list(collect_episodes(robotgym_dir)) @@ -249,11 +260,12 @@ def main(): nc.start_recording() t = time.time() dt = 1.0 / FPS - prev_ik: dict[str, float] | None = None - prev_target_ik: dict[str, float] | None = None + prev_ik: dict[str, float] | None = dict(INITIAL_JOINT_POSITIONS) + prev_target_ik: dict[str, float] | None = dict(INITIAL_JOINT_POSITIONS) for step in range(n_steps): pos_m = states[step, :3] / 10.0 # decimeters → meters + pos_m[0] += 0.1 joint_pos: dict[str, float] | None = None # IK for current EE pose. Keep prev_ik as a {joint_name: angle} dict: