diff --git a/README.md b/README.md index 3145651..872302b 100644 --- a/README.md +++ b/README.md @@ -193,6 +193,25 @@ You can also pass the number of worker processes: uv run export_mcap.py ./train_run_1 ./out 8 ``` +To derive YAM grasp-site poses from the recorded six-joint arm telemetry, opt +in with `--derive-ee-poses` on either exporter: + +```bash +uv run export_hf_task.py \ + --task organize_the_condiment_bottles \ + --derive-ee-poses + +# Or, for already-downloaded MCAPs: +uv run export_mcap.py ./train_run_1 ./out 4 --derive-ee-poses +``` + +This flag writes a separate `end_effector_poses.npz` sidecar. It does not +change `states_actions.bin`, the video, or any legacy output when the flag is +off. Poses are computed with MuJoCo from the official bundled i2rt YAM +`grasp_site` and are always expressed in each arm's **local base frame**. The +exporter intentionally does not guess a transform between the left and right +arm bases. + Each output episode is written to `./out/episode_/` in the same format the trainer reads: @@ -201,8 +220,24 @@ episode_/ states_actions.bin # (num_steps, 28) float64: 14 states + 14 actions combined_camera-images-rgb.mp4 # 30 fps vertical stack of 224x224 camera views episode_metadata.json # task name, cameras, resolutions, timing, num_steps + end_effector_poses.npz # optional; present only with --derive-ee-poses ``` +The optional NPZ uses schema `abc.end_effector_poses.v1` and contains: + +- `left_arm_state_pose`, `right_arm_state_pose`, `left_arm_action_pose`, and + `right_arm_action_pose`: `(num_steps, 7)` float64 arrays in + `[x, y, z, qw, qx, qy, qz]` order (metres, scalar-first quaternion). +- `timestamp_ns`: the aligned 30 Hz timestamps as int64 nanoseconds. +- `valid_mask`: one uint8 per step. Bits 0–3 correspond to the four pose + arrays in the order above. A bit is set only when that raw six-joint topic + supplied a finite source sample. Missing or malformed source topics produce + NaN poses with an unset bit; the legacy zero-padding is never fed to FK. + +When enabled, `episode_metadata.json` also records the sidecar schema, array +and bit assignments, arm-local frame and units, derivation provenance, bundled +model path, and SHA-256 of the exact `yam.xml` used. + The mp4 is encoded in a manner that allows for efficient dataloading. For details, see the ABC paper. ## Licenses diff --git a/abc_minimal/end_effector_poses.py b/abc_minimal/end_effector_poses.py new file mode 100644 index 0000000..6456004 --- /dev/null +++ b/abc_minimal/end_effector_poses.py @@ -0,0 +1,238 @@ +"""Derive YAM grasp-site poses from aligned ABC arm joint telemetry. + +The exporter deliberately emits these values as a separate, opt-in sidecar. +It does not change the 28-D state/action training tensor, and it does not guess +the transform between the two physical arm bases. Every pose is expressed in +the local base frame of the arm that produced the joint values. +""" + +from __future__ import annotations + +import hashlib +import os +import tempfile +from functools import lru_cache +from pathlib import Path +from typing import Final + +import mujoco +import numpy as np + +SCHEMA: Final = "abc.end_effector_poses.v1" +SIDECAR_FILENAME: Final = "end_effector_poses.npz" +POSE_FORMAT: Final = "xyz_quaternion_wxyz" +POSE_FRAME: Final = "arm_local" +SITE_NAME: Final = "grasp_site" +MODEL_RELATIVE_PATH: Final = Path("assets/put_bottles/assets/i2rt_yam/yam.xml") + +# The bit assignments are part of the v1 schema. Keep their order stable. +POSE_SPECS: Final = ( + ("left_arm_state_pose", "/left-arm-state", 0), + ("right_arm_state_pose", "/right-arm-state", 1), + ("left_arm_action_pose", "/left-arm-action", 2), + ("right_arm_action_pose", "/right-arm-action", 3), +) + + +def yam_model_path() -> Path: + """Return the official YAM MJCF bundled with this repository.""" + + return Path(__file__).resolve().parents[1] / MODEL_RELATIVE_PATH + + +@lru_cache(maxsize=1) +def yam_model_sha256() -> str: + """Return a content digest for the exact MJCF used for derivation.""" + + return hashlib.sha256(yam_model_path().read_bytes()).hexdigest() + + +class YAMForwardKinematics: + """MuJoCo forward kinematics for the bundled six-joint YAM arm chain.""" + + def __init__(self) -> None: + model_path = yam_model_path() + if not model_path.is_file(): + raise FileNotFoundError(f"bundled YAM model not found: {model_path}") + + # Loading from the file path, rather than an XML string, lets MuJoCo + # resolve the model's vendored mesh assets relative to yam.xml. + self.model = mujoco.MjModel.from_xml_path(str(model_path)) + self.data = mujoco.MjData(self.model) + + joint_ids = np.asarray( + [ + mujoco.mj_name2id( + self.model, + mujoco.mjtObj.mjOBJ_JOINT, + f"joint{joint_index}", + ) + for joint_index in range(1, 7) + ], + dtype=np.int32, + ) + if np.any(joint_ids < 0): + raise RuntimeError("bundled YAM model is missing one or more arm joints") + self._qpos_addresses = self.model.jnt_qposadr[joint_ids].copy() + + self._site_id = mujoco.mj_name2id( + self.model, + mujoco.mjtObj.mjOBJ_SITE, + SITE_NAME, + ) + if self._site_id < 0: + raise RuntimeError(f"bundled YAM model is missing {SITE_NAME}") + + def grasp_pose(self, joints: np.ndarray) -> np.ndarray: + """Return ``[x, y, z, qw, qx, qy, qz]`` in the arm-local frame.""" + + values = np.asarray(joints, dtype=np.float64) + if values.shape != (6,) or not np.all(np.isfinite(values)): + raise ValueError("joints must be six finite values in radians") + + self.data.qpos[self._qpos_addresses] = values + mujoco.mj_forward(self.model, self.data) + + quaternion = np.empty(4, dtype=np.float64) + mujoco.mju_mat2Quat(quaternion, self.data.site_xmat[self._site_id]) + quaternion /= np.linalg.norm(quaternion) + # q and -q represent the same rotation. Canonicalizing the sign keeps + # serialized results stable across MuJoCo versions. + if quaternion[0] < 0: + quaternion = -quaternion + return np.concatenate((self.data.site_xpos[self._site_id].copy(), quaternion)) + + +@lru_cache(maxsize=1) +def _kinematics() -> YAMForwardKinematics: + """Cache one compiled model and mutable data object per exporter process.""" + + return YAMForwardKinematics() + + +def _floor_indices(source_ts: np.ndarray, target_ts: np.ndarray) -> np.ndarray: + # Keep -1 for targets before the first source sample. The main exporter + # starts after every active stream, but retaining this distinction here + # prevents a future caller from silently using a sample from the future. + return np.searchsorted(source_ts, target_ts, side="right") - 1 + + +def derive_end_effector_poses( + scalars: dict[str, list[tuple[int, np.ndarray]]], + timestamp_ns: np.ndarray, +) -> dict[str, np.ndarray]: + """Derive aligned arm-local poses and a source-validity mask. + + Missing topics and malformed/non-finite joint samples remain NaN and leave + the corresponding validity bit clear. In particular, this function never + applies FK to the zero-padding used by the legacy training tensor. + """ + + ticks = np.asarray(timestamp_ns, dtype=np.int64) + if ticks.ndim != 1: + raise ValueError(f"timestamp_ns must be one-dimensional, got {ticks.shape}") + + poses = { + key: np.full((len(ticks), 7), np.nan, dtype=np.float64) + for key, _, _ in POSE_SPECS + } + valid_mask = np.zeros(len(ticks), dtype=np.uint8) + + for key, topic, bit_index in POSE_SPECS: + messages = scalars.get(topic) + if not messages: + continue + + messages = sorted(messages, key=lambda item: item[0]) + source_ts = np.asarray([time_ns for time_ns, _ in messages], dtype=np.int64) + selected = _floor_indices(source_ts, ticks) + output = poses[key] + for row_index, source_index in enumerate(selected): + if source_index < 0: + continue + joints = np.asarray(messages[int(source_index)][1], dtype=np.float64) + if joints.shape != (6,) or not np.all(np.isfinite(joints)): + continue + output[row_index] = _kinematics().grasp_pose(joints) + valid_mask[row_index] |= np.uint8(1 << bit_index) + + return {**poses, "timestamp_ns": ticks.copy(), "valid_mask": valid_mask} + + +def sidecar_metadata() -> dict: + """Return the self-describing metadata stored in episode_metadata.json.""" + + return { + "schema": SCHEMA, + "file": SIDECAR_FILENAME, + "pose_arrays": [key for key, _, _ in POSE_SPECS], + "pose_format": POSE_FORMAT, + "frame": POSE_FRAME, + "site": SITE_NAME, + "units": { + "position": "meter", + "joint_angle": "radian", + "timestamp": "nanosecond", + }, + "valid_mask_bits": { + str(bit_index): {"array": key, "source_topic": topic} + for key, topic, bit_index in POSE_SPECS + }, + "provenance": { + "method": "mujoco.mj_forward", + "mujoco_version": mujoco.__version__, + "joint_alignment": "fixed_clock_30hz_causal_floor", + "source": "raw MCAP arm joint telemetry", + }, + "model": { + "path": MODEL_RELATIVE_PATH.as_posix(), + "sha256": yam_model_sha256(), + }, + } + + +def write_end_effector_sidecar( + out_dir: Path, + scalars: dict[str, list[tuple[int, np.ndarray]]], + timestamp_ns: np.ndarray, +) -> dict: + """Write ``end_effector_poses.npz`` and return its episode metadata.""" + + arrays = derive_end_effector_poses(scalars, timestamp_ns) + out_dir = Path(out_dir) + target = out_dir / SIDECAR_FILENAME + temporary_path: Path | None = None + try: + with tempfile.NamedTemporaryFile( + mode="w+b", + dir=out_dir, + prefix=f".{SIDECAR_FILENAME}.", + suffix=".tmp", + delete=False, + ) as temporary: + temporary_path = Path(temporary.name) + np.savez_compressed(temporary, **arrays) + temporary.flush() + os.fsync(temporary.fileno()) + os.replace(temporary_path, target) + finally: + if temporary_path is not None: + temporary_path.unlink(missing_ok=True) + return sidecar_metadata() + + +__all__ = [ + "MODEL_RELATIVE_PATH", + "POSE_FORMAT", + "POSE_FRAME", + "POSE_SPECS", + "SCHEMA", + "SIDECAR_FILENAME", + "SITE_NAME", + "YAMForwardKinematics", + "derive_end_effector_poses", + "sidecar_metadata", + "write_end_effector_sidecar", + "yam_model_path", + "yam_model_sha256", +] diff --git a/export_hf_task.py b/export_hf_task.py index 76bb13f..a4a5662 100644 --- a/export_hf_task.py +++ b/export_hf_task.py @@ -1,6 +1,6 @@ # /// script # requires-python = ">=3.10" -# dependencies = ["numpy", "mcap", "mcap-protobuf-support", "tyro"] +# dependencies = ["numpy", "mcap", "mcap-protobuf-support", "mujoco", "tyro"] # /// """Export one ABC-130k Hugging Face task to the training format.""" @@ -43,6 +43,12 @@ class Config: max_episodes: Annotated[int | None, tyro.conf.arg(help="Optional per-split cap for smoke tests.")] = None dry_run: Annotated[bool, tyro.conf.arg(help="List only; do not download or convert.")] = False keep_mcaps: Annotated[bool, tyro.conf.arg(help="Keep staged raw MCAPs after conversion.")] = False + derive_ee_poses: Annotated[ + bool, + tyro.conf.arg( + help="Write arm-local YAM grasp poses to end_effector_poses.npz." + ), + ] = False def token(cfg: Config) -> str | None: @@ -167,6 +173,8 @@ def write_manifest(cfg: Config, split: str, files: list[dict], root: Path) -> No def convert(cfg: Config, split: str, root: Path) -> None: cmd = [sys.executable, "export_mcap.py", str(root), str((cfg.cache / f"{split}_real").expanduser()), str(cfg.workers)] + if cfg.derive_ee_poses: + cmd.append("--derive-ee-poses") print("[convert]", " ".join(cmd)) subprocess.run(cmd, check=True) diff --git a/export_mcap.py b/export_mcap.py index 020cab2..8af9d66 100644 --- a/export_mcap.py +++ b/export_mcap.py @@ -1,6 +1,6 @@ # /// script # requires-python = ">=3.10" -# dependencies = ["numpy", "mcap", "mcap-protobuf-support", "tyro"] +# dependencies = ["numpy", "mcap", "mcap-protobuf-support", "mujoco", "tyro"] # /// """Convert release-format MCAP episodes into the training data layout. @@ -63,6 +63,12 @@ class ExportMcapConfig: root: Annotated[Path, tyro.conf.Positional] out_dir: Annotated[Path, tyro.conf.Positional] workers: Annotated[int, tyro.conf.Positional] = 4 + derive_ee_poses: Annotated[ + bool, + tyro.conf.arg( + help="Write arm-local YAM grasp poses to end_effector_poses.npz." + ), + ] = False def floor_indices(source_ts, target_ts): @@ -110,8 +116,18 @@ def encode_aligned(h264_path, width, height, needed, out_path): raise RuntimeError("ffmpeg encode failed") +def _unpack_export_job(job): + """Accept current jobs and the historical three-item direct-call form.""" + if len(job) == 3: + mcap_path, task_name, out_root = job + return mcap_path, task_name, out_root, False + if len(job) == 4: + return job + raise ValueError(f"export job must contain 3 or 4 items, got {len(job)}") + + def export_episode(job): - mcap_path, task_name, out_root = job + mcap_path, task_name, out_root, derive_ee_poses = _unpack_export_job(job) from mcap.reader import make_reader from mcap_protobuf.decoder import DecoderFactory @@ -213,6 +229,10 @@ def export_episode(job): "camera_resolutions": {k: [OUT_W, OUT_H] for k, _ in active_cams}, "alignment": "fixed_clock_30hz_causal", "t0_ns": int(t0), "tick_ns": TICK_NS, "num_steps": num_steps} + if derive_ee_poses: + from abc_minimal.end_effector_poses import write_end_effector_sidecar + + meta["end_effector_poses"] = write_end_effector_sidecar(out_dir, scalars, ticks) (out_dir / "episode_metadata.json").write_text(json.dumps(meta, indent=2)) print(f"[OK] {ep_id}: {num_steps} steps, cams={[k for k, _ in active_cams]}") return ep_id @@ -220,7 +240,7 @@ def export_episode(job): def main(config: ExportMcapConfig): jobs = sorted( - (str(p), p.parent.parent.name, str(config.out_dir)) + (str(p), p.parent.parent.name, str(config.out_dir), config.derive_ee_poses) for p in config.root.glob("*/episode_*/episode.mcap") ) print(f"{len(jobs)} episodes") diff --git a/pyproject.toml b/pyproject.toml index 16e1e54..35035ea 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,6 +26,11 @@ dependencies = [ "wandb", ] +[dependency-groups] +dev = [ + "pytest>=8", +] + [tool.uv.sources] torch = { index = "pytorch-cu128" } torchcodec = { index = "pytorch-cpu" } diff --git a/tests/test_end_effector_poses.py b/tests/test_end_effector_poses.py new file mode 100644 index 0000000..18e2341 --- /dev/null +++ b/tests/test_end_effector_poses.py @@ -0,0 +1,201 @@ +"""Regression tests for the opt-in end-effector pose sidecar.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import mujoco +import numpy as np + +from abc_minimal.end_effector_poses import ( + POSE_SPECS, + SCHEMA, + SIDECAR_FILENAME, + YAMForwardKinematics, + derive_end_effector_poses, + sidecar_metadata, + write_end_effector_sidecar, + yam_model_path, + yam_model_sha256, +) + +# Literal telemetry and transforms from the first arm-state messages in this +# ungated public ABC mirror (sequence 0, 2025-04-03T17:28:55.850143744Z): +# https://huggingface.co/datasets/Voxel51/ABC-130k/blob/main/data/val/ +# fold_and_stack_the_skirts/episode_aa015e3a-8cb9-4e03-93a9-900fac9ae6b8/ +# episode.fo.mcap +# Mirror revision: 9659e8ce4b39580f48369cc31bc2e47a217c40e7. +LEFT_JOINTS = np.array( + [ + -0.48161287861448265, + 0.7116426337071786, + 0.7040131227588304, + -0.7028686961165764, + -0.012016479743651942, + -0.13103685053788006, + ] +) +LEFT_ARM_LOCAL_TRANSFORM = np.array( + [ + [ + -0.632531141717365, + 0.3756759502397681, + 0.6773270518510042, + 0.2229318922544835, + ], + [ + 0.18317820509999394, + 0.9222451495056406, + -0.3404550328452526, + -0.11421321941931482, + ], + [ + -0.7525623561765962, + -0.09127685700529321, + -0.6521644236242606, + 0.1754101887061787, + ], + [0.0, 0.0, 0.0, 1.0], + ] +) +RIGHT_JOINTS = np.array( + [ + 0.43926909285114846, + 1.0530632486457616, + 0.7162203402761875, + -0.7997634851606019, + -0.03337911039902686, + 0.0997558556496525, + ] +) +# The released right-arm record is in a shared bimanual frame. It differs +# from the arm-local model output only by an observed -0.61 m Y translation. +# This fixture removes that translation; production code never applies or +# records the disputed shared-frame transform. +RIGHT_ARM_LOCAL_TRANSFORM = np.array( + [ + [ + -0.8605964623532097, + -0.35379918654529247, + 0.3663329968553596, + 0.21752706566379915, + ], + [ + -0.2944080765331712, + 0.9325444245486542, + 0.20900904457763797, + 0.1084818535402495, + ], + [ + -0.41556902369784254, + 0.0720210513885209, + -0.906705770743582, + 0.08778088368163128, + ], + [0.0, 0.0, 0.0, 1.0], + ] +) + + +def pose_matrix(pose: np.ndarray) -> np.ndarray: + matrix = np.eye(4, dtype=np.float64) + rotation = np.empty(9, dtype=np.float64) + mujoco.mju_quat2Mat(rotation, pose[3:]) + matrix[:3, :3] = rotation.reshape(3, 3) + matrix[:3, 3] = pose[:3] + return matrix + + +def test_official_model_matches_public_left_and_right_arm_local_fixtures() -> None: + fk = YAMForwardKinematics() + + np.testing.assert_allclose( + pose_matrix(fk.grasp_pose(LEFT_JOINTS)), + LEFT_ARM_LOCAL_TRANSFORM, + rtol=0.0, + atol=1e-12, + ) + np.testing.assert_allclose( + pose_matrix(fk.grasp_pose(RIGHT_JOINTS)), + RIGHT_ARM_LOCAL_TRANSFORM, + rtol=0.0, + atol=1e-12, + ) + + +def test_missing_topics_are_nan_and_never_marked_valid() -> None: + ticks = np.array([90, 110, 120], dtype=np.int64) + scalars = { + "/left-arm-state": [(100, LEFT_JOINTS)], + # Present but malformed is invalid just like an absent source sample. + "/right-arm-action": [(100, np.zeros(5))], + } + + result = derive_end_effector_poses(scalars, ticks) + + assert result["timestamp_ns"].dtype == np.int64 + assert result["valid_mask"].dtype == np.uint8 + np.testing.assert_array_equal( + result["valid_mask"], + np.array([0, 1, 1], dtype=np.uint8), + ) + assert np.isnan(result["left_arm_state_pose"][0]).all() + assert np.isfinite(result["left_arm_state_pose"][1:]).all() + for key in ( + "right_arm_state_pose", + "left_arm_action_pose", + "right_arm_action_pose", + ): + assert np.isnan(result[key]).all() + + +def test_all_four_validity_bits_and_causal_alignment() -> None: + ticks = np.array([101, 201], dtype=np.int64) + scalars = { + topic: [(100, LEFT_JOINTS), (200, RIGHT_JOINTS)] for _, topic, _ in POSE_SPECS + } + + result = derive_end_effector_poses(scalars, ticks) + + np.testing.assert_array_equal( + result["valid_mask"], np.array([15, 15], dtype=np.uint8) + ) + for key, _, _ in POSE_SPECS: + np.testing.assert_allclose( + result[key][0], YAMForwardKinematics().grasp_pose(LEFT_JOINTS) + ) + np.testing.assert_allclose( + result[key][1], YAMForwardKinematics().grasp_pose(RIGHT_JOINTS) + ) + + +def test_sidecar_and_episode_metadata_are_self_describing(tmp_path: Path) -> None: + ticks = np.array([101], dtype=np.int64) + # An interrupted or repeated export may already have a target file. The + # writer must atomically replace it and leave no temporary artifact. + (tmp_path / SIDECAR_FILENAME).write_bytes(b"stale") + metadata = write_end_effector_sidecar( + tmp_path, + {"/left-arm-state": [(100, LEFT_JOINTS)]}, + ticks, + ) + + with np.load(tmp_path / SIDECAR_FILENAME) as sidecar: + assert set(sidecar.files) == { + *(key for key, _, _ in POSE_SPECS), + "timestamp_ns", + "valid_mask", + } + assert sidecar["left_arm_state_pose"].shape == (1, 7) + assert {path.name for path in tmp_path.iterdir()} == {SIDECAR_FILENAME} + assert metadata == sidecar_metadata() + assert metadata["schema"] == SCHEMA + assert metadata["frame"] == "arm_local" + assert metadata["provenance"]["mujoco_version"] == mujoco.__version__ + assert metadata["model"]["sha256"] == yam_model_sha256() + assert len(metadata["model"]["sha256"]) == 64 + assert yam_model_path().is_file() + # Ensure the descriptor remains JSON serializable before the exporter + # inserts it into episode_metadata.json. + json.dumps(metadata) diff --git a/tests/test_export_flags.py b/tests/test_export_flags.py new file mode 100644 index 0000000..f9d35c5 --- /dev/null +++ b/tests/test_export_flags.py @@ -0,0 +1,53 @@ +"""CLI plumbing tests for the opt-in geometry export.""" + +from __future__ import annotations + +from pathlib import Path + +import export_hf_task +from export_mcap import ExportMcapConfig, _unpack_export_job + + +def test_geometry_export_is_off_by_default(tmp_path: Path) -> None: + config = ExportMcapConfig(root=tmp_path, out_dir=tmp_path) + + assert config.derive_ee_poses is False + + +def test_historical_three_item_export_jobs_remain_compatible() -> None: + historical = ("episode.mcap", "task", "out") + current = (*historical, True) + + assert _unpack_export_job(historical) == (*historical, False) + assert _unpack_export_job(current) == current + + +def test_hf_wrapper_only_forwards_explicit_geometry_flag( + tmp_path: Path, + monkeypatch, +) -> None: + commands: list[list[str]] = [] + + def capture(command, *, check): + assert check is True + commands.append(command) + + monkeypatch.setattr(export_hf_task.subprocess, "run", capture) + + export_hf_task.convert( + export_hf_task.Config(task="example", cache=tmp_path), + "train", + tmp_path / "staged", + ) + export_hf_task.convert( + export_hf_task.Config( + task="example", + cache=tmp_path, + derive_ee_poses=True, + ), + "train", + tmp_path / "staged", + ) + + assert "--derive-ee-poses" not in commands[0] + assert commands[1][-1] == "--derive-ee-poses"