feat(data): optionally export arm-local YAM end-effector poses

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_<uuid>/` in the same format
 the trainer reads:
 
@@ -201,8 +220,24 @@ episode_<uuid>/
   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

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",
+]

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)
 

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,14 +229,18 @@ 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
 
 
 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")

pyproject.toml

@@ -26,6 +26,11 @@ dependencies = [
     "wandb",
 ]
 
+[dependency-groups]
+dev = [
+    "pytest>=8",
+]
+
 [tool.uv.sources]
 torch = { index = "pytorch-cu128" }
 torchcodec = { index = "pytorch-cpu" }