From 1c03237eee171e7f675dafaf2b666d974e4a83de Mon Sep 17 00:00:00 2001 From: cc Date: Fri, 30 Oct 2026 17:31:19 -0700 Subject: [PATCH] feat(imitation): integrate native OpenYAM Quest collection (#4343) --- dimos/imitation/collection/README.md | 24 +- dimos/imitation/dataprep/README.md | 2 +- dimos/imitation/dataprep/test_mcap_source.py | 10 +- dimos/robot/all_blueprints.py | 1 + .../blueprints/learning_quest_collection.py | 76 ++++++ .../test_learning_collection_e2e.py | 234 ++++++++++++++++++ experimental/gateway/annotations.yaml | 3 + 7 files changed, 336 insertions(+), 14 deletions(-) create mode 100644 dimos/robot/manipulators/openyam/blueprints/learning_quest_collection.py create mode 100644 dimos/robot/manipulators/openyam/blueprints/test_learning_collection_e2e.py diff --git a/dimos/imitation/collection/README.md b/dimos/imitation/collection/README.md index 5a21cdaf2e4..ef56cf8ae7d 100644 --- a/dimos/imitation/collection/README.md +++ b/dimos/imitation/collection/README.md @@ -6,7 +6,7 @@ controls through ordinary `autoconnect()` and `remappings()`. ## Declare the recorded features -This example declares measured joint state, accepted joint targets, and one RGB +This example declares measured joint state, commanded joint targets, and one RGB image. Constructing the profile does not connect to hardware. ```python session=profile no-result @@ -29,13 +29,21 @@ profile = CollectionProfile( names=["height", "width", "channels"], ), }, - actions={"action": state.model_copy(update={ - "stream": "accepted", "source_kind": "joint_position_updates", - }, deep=True)}, + actions={ + "action": state.model_copy( + update={ + "stream": "commanded", + "source_kind": "joint_position_updates", + }, + deep=True, + ) + }, metadata={"name": "custom-arm", "robot_type": "custom-arm-2dof"}, ) assert profile.input_types() == { - "measured": JointState, "wrist_image": Image, "accepted": JointState, + "measured": JointState, + "wrist_image": Image, + "commanded": JointState, } ``` @@ -66,7 +74,7 @@ recorder = collect(profile=profile, recording=Path("recordings/session-001")) Include `recorder`, your robot, camera producers, and an `EpisodeMonitorModule.blueprint(instance_name="episodes")` in `autoconnect()`. The monitor supplies `status: In[String]` containing versioned JSON episode events. -Robot-specific readiness rules belong on that monitor. The recorder must precede +The monitor only controls episode start, save, and discard. The recorder must precede producers in the Blueprint so it stops after them. The Python graph supports any number of coordinators. Connect each recorded stream @@ -131,8 +139,8 @@ add an explicit depth feature and connect its producer. Capture stores asynchronous streams at their native rates. The saved schema fixes the alignment to a 30 Hz grid, anchored at the first video observation, or the first observation when there is no video. Snapshot matching uses a 20 ms tolerance. -Accepted joint updates use causal hold: only commands at or before the target time -contribute, and omitted joints retain their last accepted target. Missing accepted +Commanded joint updates use causal hold: only commands at or before the target time +contribute, and omitted joints retain their last commanded target. Missing commanded targets never fall back to measured state or future commands. The collection API exposes no `sync` or `quality` options. Existing low-level diff --git a/dimos/imitation/dataprep/README.md b/dimos/imitation/dataprep/README.md index 3706c3cdb79..e0c8e38293c 100644 --- a/dimos/imitation/dataprep/README.md +++ b/dimos/imitation/dataprep/README.md @@ -85,7 +85,7 @@ config = DataPrepConfig( }, action={ "action": FeatureSpec( - stream="applied_joint_position_command", + stream="coordinator_joint_command", field="position", dtype="float32", shape=(2,), diff --git a/dimos/imitation/dataprep/test_mcap_source.py b/dimos/imitation/dataprep/test_mcap_source.py index 5483d6915b3..e9b00ec40a5 100644 --- a/dimos/imitation/dataprep/test_mcap_source.py +++ b/dimos/imitation/dataprep/test_mcap_source.py @@ -82,8 +82,8 @@ def _write_collection(path: Path) -> None: "coordinator_joint_state": _register_channel( writer, "coordinator_joint_state", JointState ), - "applied_joint_position_command": _register_channel( - writer, "applied_joint_position_command", JointState + "coordinator_joint_command": _register_channel( + writer, "coordinator_joint_command", JointState ), "status": _register_channel(writer, "status", String, "json"), } @@ -130,7 +130,7 @@ def _write_collection(path: Path) -> None: ) _write_message( writer, - channels["applied_joint_position_command"], + channels["coordinator_joint_command"], ts, JointState( ts=ts, @@ -179,7 +179,7 @@ def _config(source: Path, output: Path) -> DataPrepConfig: }, action={ "action": FeatureSpec( - stream="applied_joint_position_command", + stream="coordinator_joint_command", field="position", dtype="float32", shape=(2,), @@ -213,7 +213,7 @@ def writer(samples: Iterator[Sample], selected_output: OutputConfig) -> Path: dataset_path = run_dataprep(config, writer=writer) assert info["streams"] == { - "applied_joint_position_command": 3, + "coordinator_joint_command": 3, "color_image": 3, "coordinator_joint_state": 3, "status": 2, diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index cd285b95ba8..568ee291b5b 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -109,6 +109,7 @@ "mid360-realsense-record-with-pcap": "dimos.robot.assembly.mid360_realsense_30:mid360_realsense_record_with_pcap", "openarm-planner-coordinator": "dimos.robot.manipulators.openarm.blueprints.basic:openarm_planner_coordinator", "openyam-planner-coordinator": "dimos.robot.manipulators.openyam.blueprints.basic:openyam_planner_coordinator", + "openyam-quest-collection": "dimos.robot.manipulators.openyam.blueprints.learning_quest_collection:openyam_quest_collection", "r1pro-coordinator": "dimos.robot.galaxea.r1pro.blueprints.basic.r1pro_coordinator:r1pro_coordinator", "r1pro-manipulation": "dimos.robot.galaxea.r1pro.blueprints.manipulation.r1pro_manipulation:r1pro_manipulation", "r1pro-nav": "dimos.robot.galaxea.r1pro.blueprints.navigation.r1pro_nav:r1pro_nav", diff --git a/dimos/robot/manipulators/openyam/blueprints/learning_quest_collection.py b/dimos/robot/manipulators/openyam/blueprints/learning_quest_collection.py new file mode 100644 index 00000000000..9aa6d05e8d0 --- /dev/null +++ b/dimos/robot/manipulators/openyam/blueprints/learning_quest_collection.py @@ -0,0 +1,76 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Quest-guided OpenYAM collection.""" + +from dimos.core.coordination.blueprints import autoconnect +from dimos.hardware.sensors.camera.module import WebcamModule +from dimos.hardware.sensors.camera.webcam import WebcamConfig +from dimos.imitation.collection.episode_monitor import EpisodeMonitorModule +from dimos.imitation.collection.profile import CollectionProfile +from dimos.imitation.collection.recorder import collect +from dimos.imitation.dataprep.schema import FeatureSpec +from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.robot.manipulators.openyam.blueprints.teleop import teleop_webxr_openyam +from dimos.robot.manipulators.openyam.joints import OPENYAM_JOINTS + +_state = FeatureSpec( + stream="coordinator_joint_state", + message_type=JointState, + field="position", + dtype="float32", + shape=(len(OPENYAM_JOINTS),), + names=list(OPENYAM_JOINTS), +) +_profile = CollectionProfile( + observations={ + "observation.images.wrist": FeatureSpec( + stream="wrist_image", + message_type=Image, + field="data", + dtype="video", + shape=(480, 640, 3), + names=["height", "width", "channels"], + ), + "observation.state": _state, + }, + actions={ + "action": _state.model_copy( + update={ + "stream": "coordinator_joint_command", + "source_kind": "joint_position_updates", + }, + deep=True, + ) + }, + metadata={"name": "openyam-quest", "robot_type": "openyam"}, +) + +openyam_quest_collection = autoconnect( + collect(profile=_profile), + EpisodeMonitorModule.blueprint(instance_name="episodes"), + teleop_webxr_openyam, + WebcamModule.blueprint( + instance_name="wrist", + hardware=WebcamConfig(camera_index=0, fps=30, frame_id_prefix="wrist_image"), + frame_id="wrist_camera_link", + ), +).remappings( + [ + ("wrist", "color_image", "wrist_image"), + ("wrist", "camera_info", "wrist_camera_info"), + ("wrist", "tf", "wrist_tf"), + ] +) diff --git a/dimos/robot/manipulators/openyam/blueprints/test_learning_collection_e2e.py b/dimos/robot/manipulators/openyam/blueprints/test_learning_collection_e2e.py new file mode 100644 index 00000000000..ee79ba9e996 --- /dev/null +++ b/dimos/robot/manipulators/openyam/blueprints/test_learning_collection_e2e.py @@ -0,0 +1,234 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Native OpenYAM collection interoperability over the production transport.""" + +from __future__ import annotations + +import os +from pathlib import Path +import subprocess +from typing import Any, cast +import uuid + +import numpy as np +import pytest + +from dimos.constants import DIMOS_PROJECT_ROOT +from dimos.core.global_config import global_config +from dimos.core.transport import ZenohTransport +from dimos.imitation.collection.episode import EpisodeStatus +from dimos.imitation.dataprep.core import extract_episodes +from dimos.imitation.dataprep.schema import EpisodeExtractor +from dimos.memory.store.sqlite import SqliteStore +from dimos.msgs.protocol import DimosMsg +from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.msgs.std_msgs.String import String +from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.protocol.pubsub.impl.zenohpubsub import QOS_NEVER_DROP, Topic +from dimos.robot.manipulators.openyam.blueprints.learning_quest_collection import ( + openyam_quest_collection, +) +from dimos.utils.testing.waiting import wait_until + +pytestmark = [ + pytest.mark.native_e2e, + pytest.mark.skipif_macos, + pytest.mark.skipif_aarch64, + pytest.mark.skipif_no_turbojpeg, +] + +_RUST_WORKSPACE = DIMOS_PROJECT_ROOT / "dimos" / "experimental" / "memory" / "rust" +_EXECUTABLE = Path( + os.environ.get( + "DIMOS_MEMORY_RECORDER_TEST_EXECUTABLE", + str(DIMOS_PROJECT_ROOT / "target" / "debug" / "dimos-memory-recorder"), + ) +) + + +@pytest.fixture(scope="module") +def native_recorder_executable() -> Path: + if "DIMOS_MEMORY_RECORDER_TEST_EXECUTABLE" not in os.environ: + subprocess.run( + ["cargo", "build", "--locked", "-p", "dimos-memory-recorder"], + cwd=_RUST_WORKSPACE, + check=True, + ) + return _EXECUTABLE + + +def _stream_count(path: Path, name: str) -> int: + with SqliteStore(path=str(path), must_exist=True) as store: + return store.stream(name).count() + + +def test_native_collection_records_typed_zenoh_streams( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + native_recorder_executable: Path, +) -> None: + monkeypatch.setattr(global_config, "transport", "zenoh") + directory = tmp_path / "native-openyam" + artifact = directory / "recording.db" + atom = next( + atom for atom in openyam_quest_collection.active_blueprints if atom.name == "recorder" + ) + recorder = atom.module( + executable=str(native_recorder_executable), + source_dir=None, + build_command=None, + **{**atom.kwargs, "recording": directory, "format": "sqlite"}, + ) + topic_prefix = f"dimos/test/native-collection/{uuid.uuid4().hex}" + payload_types: dict[str, type[Any]] = { + "wrist_image": Image, + "coordinator_joint_state": JointState, + "coordinator_joint_command": JointState, + "status": String, + "tf": TFMessage, + } + publishers: dict[str, ZenohTransport[Any]] = {} + for name, payload_type in payload_types.items(): + transport = ZenohTransport( + Topic( + f"{topic_prefix}/{name}", + cast("type[DimosMsg]", payload_type), + qos=QOS_NEVER_DROP, + ) + ) + getattr(recorder, name).transport = transport + publishers[name] = transport + + def publish(name: str, message: Any) -> None: + publishers[name].broadcast( + None, String(message.to_json()) if isinstance(message, EpisodeStatus) else message + ) + + try: + recorder.start() + ready = EpisodeStatus( + ts=1.0, + state="idle", + last_event="init", + episodes_saved=0, + episodes_discarded=0, + ) + wait_until( + lambda: (publish("status", ready), _stream_count(artifact, "status") > 0)[1], + timeout=10.0, + interval=0.1, + message="native recorder did not receive the Zenoh status stream", + ) + + publish( + "status", + EpisodeStatus( + ts=10.0, + state="recording", + last_event="start", + episodes_saved=0, + episodes_discarded=0, + task_label="pick", + ), + ) + publish( + "wrist_image", + Image( + ts=10.0, + frame_id="wrist_camera_link", + format=ImageFormat.RGB, + data=np.full((16, 16, 3), 127, dtype=np.uint8), + ), + ) + publish( + "coordinator_joint_state", + JointState( + ts=10.0, + frame_id="coordinator", + name=["shoulder", "wrist"], + position=[0.25, -0.5], + velocity=[0.0, 0.0], + effort=[0.0, 0.0], + ), + ) + publish( + "coordinator_joint_command", + JointState( + ts=10.0, + frame_id="coordinator", + name=["shoulder", "wrist"], + position=[0.3, -0.4], + velocity=[], + effort=[], + ), + ) + publish( + "status", + EpisodeStatus( + ts=11.0, + state="idle", + last_event="save", + episodes_saved=1, + episodes_discarded=0, + task_label="pick", + ), + ) + wait_until( + lambda: all( + _stream_count(artifact, name) >= count + for name, count in { + "wrist_image": 1, + "coordinator_joint_state": 1, + "coordinator_joint_command": 1, + "status": 3, + }.items() + ), + timeout=10.0, + interval=0.05, + message="native recorder did not persist every collection stream", + ) + finally: + recorder.stop() + for publisher in publishers.values(): + publisher.stop() + + with SqliteStore(path=str(artifact), must_exist=True) as store: + image_observation = store.stream("wrist_image", Image).last() + image = image_observation.data + image_ts = image_observation.ts + joint_observation = store.stream("coordinator_joint_state", JointState).last() + joints = joint_observation.data + joint_ts = joint_observation.ts + action = store.stream("coordinator_joint_command", JointState).last().data + statuses = [ + (observation.ts, EpisodeStatus.from_json(observation.data.data).last_event) + for observation in store.stream("status", String).to_list() + ] + episodes = extract_episodes(store, EpisodeExtractor(status_stream="status")) + + assert image_ts == 10.0 + assert image.frame_id == "wrist_camera_link" + assert image.data.shape == (16, 16, 3) + assert joint_ts == 10.0 + assert joints.position == [0.25, -0.5] + assert action.position == [0.3, -0.4] + assert statuses[-2:] == [ + (10.0, "start"), + (11.0, "save"), + ] + assert [(episode.start_ts, episode.end_ts, episode.task_label) for episode in episodes] == [ + (10.0, 11.0, "pick") + ] diff --git a/experimental/gateway/annotations.yaml b/experimental/gateway/annotations.yaml index 9d9f80e2bae..82b7f22e091 100644 --- a/experimental/gateway/annotations.yaml +++ b/experimental/gateway/annotations.yaml @@ -631,6 +631,9 @@ robots: coordinator-openyam: title: OpenYAM coordinator description: An OpenYAM arm behind a ControlCoordinator (a building block). + openyam-quest-collection: + title: OpenYAM demonstration recording + description: Teleoperate an OpenYAM arm in VR and record demonstrations. teleop-webxr-openyam: title: OpenYAM VR teleop description: Move an OpenYAM arm with a VR headset's right controller.