diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index fc88b3763f..a9f8291eec 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -329,10 +329,10 @@ jobs: with: enable-cache: true prune-cache: true - - name: Install dependency for pyaudio + - name: Install Python audio and image dependencies run: | sudo apt-get update - sudo apt-get install -y portaudio19-dev + sudo apt-get install -y portaudio19-dev libturbojpeg - name: Build and test PyO3 bindings # PyO3 extension modules (cdylib + `extension-module`) are the one case # that still needs a plain build: pytest imports the compiled .so. @@ -352,7 +352,10 @@ jobs: - name: Native module e2e tests run: | cargo build --release --locked -p dimos-livox -p dimos-virtual-mid360 - uv run pytest -m native_e2e dimos/hardware/sensors/lidar/livox/test_e2e.py --no-cov + uv run pytest -m native_e2e \ + dimos/hardware/sensors/lidar/livox/test_e2e.py \ + dimos/robot/manipulators/openyam/blueprints/test_learning_collection_e2e.py \ + dimos/imitation/test_datacollection_e2e.py --no-cov native: name: Native builds (C++ and Rust) diff --git a/Cargo.lock b/Cargo.lock index 5aebcae77f..03f1f0fe8a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -650,7 +650,7 @@ dependencies = [ [[package]] name = "dimos-lcm" version = "0.1.0" -source = "git+https://github.com/dimensionalOS/dimos-lcm.git?branch=rust-codegen#04d78e8622500244123ba9cefa4c51b4cb454549" +source = "git+https://github.com/dimensionalOS/dimos-lcm.git?rev=dd2159513ebfaa7ebc5fc32bf60209cd09aa1ca5#dd2159513ebfaa7ebc5fc32bf60209cd09aa1ca5" dependencies = [ "byteorder", "socket2 0.5.10", @@ -750,7 +750,7 @@ dependencies = [ "proc-macro2", "quote", "syn 3.0.4", - "toml 1.1.4+spec-1.1.0", + "toml 1.1.5+spec-1.1.0", ] [[package]] @@ -1675,7 +1675,7 @@ checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" [[package]] name = "lcm-msgs" version = "0.1.0" -source = "git+https://github.com/dimensionalOS/dimos-lcm.git?branch=rust-codegen#04d78e8622500244123ba9cefa4c51b4cb454549" +source = "git+https://github.com/dimensionalOS/dimos-lcm.git?rev=dd2159513ebfaa7ebc5fc32bf60209cd09aa1ca5#dd2159513ebfaa7ebc5fc32bf60209cd09aa1ca5" dependencies = [ "byteorder", ] @@ -1704,9 +1704,9 @@ checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" [[package]] name = "libredox" -version = "0.1.21" +version = "0.1.23" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d7955dfc218a8afb29dfeffd540e3a6e96baeb94fe7138228dd7cc6937fbbf96" +checksum = "8d8f1ea3f21fd3405dcaf6c9b5c1630af9afc422d9073ea39c5f6d6c772e08ed" dependencies = [ "libc", ] @@ -1852,9 +1852,9 @@ dependencies = [ [[package]] name = "mio" -version = "1.2.2" +version = "1.2.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "30d65c71f1ce40ab09135ce117d742b9f8a19ff91a41a8b57ed50bc2de59c427" +checksum = "4b18443e9c262bfe8fa82f51666e2642c53393f7e5c27b3e1aeab922cff5b9d8" dependencies = [ "libc", "wasi", @@ -2395,13 +2395,13 @@ dependencies = [ [[package]] name = "prometheus-client-derive-encode" -version = "0.5.0" +version = "0.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9adf1691c04c0a5ff46ff8f262b58beb07b0dbb61f96f9f54f6cbd82106ed87f" +checksum = "01e34894696ff94f64a20c2c373a6440903e9c2789a303d68ec6e6f953f890e4" dependencies = [ "proc-macro2", "quote", - "syn 2.0.119", + "syn 3.0.4", ] [[package]] @@ -3242,9 +3242,9 @@ checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" [[package]] name = "smallvec" -version = "1.15.2" +version = "1.16.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8ed6a63f02c8539c91a8685a86f4099661ba3da017932f6ebbea6de3f0fa7c90" +checksum = "b9be42f50aa861c555654aa3a37f52f4b1074bacf4e48fe0ef7fa584e80f1f0f" [[package]] name = "socket2" @@ -3553,9 +3553,9 @@ dependencies = [ [[package]] name = "tinyvec" -version = "1.12.0" +version = "1.13.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bb4ebadaa0af04fab11ae01eb5f9fdb5f9c5b875506e210e71c07873528baa7f" +checksum = "4cf0ded5c4e56918d8f8a339e1bb67d038d3bc6d144ac407904015ba2e4cde9b" dependencies = [ "tinyvec_macros", ] @@ -3635,9 +3635,9 @@ dependencies = [ [[package]] name = "toml" -version = "1.1.4+spec-1.1.0" +version = "1.1.5+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3aace63f4bbcdfc2c965b059de67119c89c4017a70d633be6c104910f67056f5" +checksum = "12c0ba9680044b4ce98d391a62094047eada0d64860b80166c39f4a6b5640785" dependencies = [ "indexmap 2.14.1", "serde_core", diff --git a/dimos/agents/demo_agent.py b/dimos/agents/demo_agent.py index 1987565369..1714554b37 100644 --- a/dimos/agents/demo_agent.py +++ b/dimos/agents/demo_agent.py @@ -17,25 +17,21 @@ from dimos.agents.skills.observe_skill import ObserveSkill from dimos.core.coordination.blueprints import autoconnect from dimos.hardware.sensors.camera.module import CameraModule -from dimos.hardware.sensors.camera.webcam import Webcam +from dimos.hardware.sensors.camera.webcam import WebcamConfig from dimos.hardware.sensors.camera.zed import compat as zed demo_agent = autoconnect(McpServer.blueprint(), McpClient.blueprint()) -def _create_webcam() -> Webcam: - return Webcam( - camera_index=0, - fps=15, - camera_info=zed.CameraInfo.SingleWebcam, - ) - - demo_agent_camera = autoconnect( McpServer.blueprint(), McpClient.blueprint(), ObserveSkill.blueprint(), CameraModule.blueprint( - hardware=_create_webcam, + hardware=WebcamConfig( + camera_index=0, + fps=15, + camera_info=zed.CameraInfo.SingleWebcam, + ), ), ) diff --git a/dimos/cli/commands/dataprep.py b/dimos/cli/commands/dataprep.py deleted file mode 100644 index 4df15b3c91..0000000000 --- a/dimos/cli/commands/dataprep.py +++ /dev/null @@ -1,52 +0,0 @@ -# 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. - -"""`dimos dataprep` commands; the implementation lives in dimos.imitation.dataprep.cli.""" - -from __future__ import annotations - -from pathlib import Path -from typing import Literal, cast - -import typer - -from dimos.imitation.dataprep.cli import build, inspect - -dataprep_app = typer.Typer(help="Build and inspect learning datasets from recordings") - - -@dataprep_app.command("build") -def dataprep_build( - source: Path | None = typer.Option(None, "--source", "-s", help="Recording .db to read"), - output: Path | None = typer.Option(None, "--output", help="Dataset output directory"), - output_format: str = typer.Option(None, "--format", "-f", help="Output format: lerobot | hdf5"), - config_path: Path | None = typer.Option( - None, "--config", "-c", help="JSON DataPrepConfig (needed for obs/action stream maps)" - ), -) -> None: - """Build a dataset from a recording (lerobot/hdf5 + dimos_meta.json).""" - build(config_path, source, output, cast("Literal['lerobot', 'hdf5'] | None", output_format)) - - -@dataprep_app.command("inspect") -def dataprep_inspect( - dataset: Path | None = typer.Argument( - None, help="Recording .db, built .hdf5 file, or lerobot directory" - ), - output_format: str = typer.Option( - None, "--format", "-f", help="lerobot | hdf5 (auto-detected from the path if omitted)" - ), -) -> None: - """Summarize a recording or built dataset, including incomplete episodes.""" - inspect(dataset, cast("Literal['lerobot', 'hdf5'] | None", output_format)) diff --git a/dimos/cli/dimos.py b/dimos/cli/dimos.py index 06899c059b..5ca2f20417 100644 --- a/dimos/cli/dimos.py +++ b/dimos/cli/dimos.py @@ -53,7 +53,6 @@ from dimos.cli.commands.bake import bake from dimos.cli.commands.cameracalibrate import cameracalibrate from dimos.cli.commands.data import data_app -from dimos.cli.commands.dataprep import dataprep_app from dimos.cli.commands.docs import docs from dimos.cli.commands.global_options import create_dynamic_callback from dimos.cli.commands.graph import graph @@ -143,7 +142,6 @@ def cli_main() -> None: from dimos.navigation.nav_3d.evaluator.cli import app as nav_eval_app main.add_typer(nav_eval_app, name="nav-eval") -main.add_typer(dataprep_app, name="dataprep") from dimos.memory.cli.app import mem_app diff --git a/dimos/control/coordinator.py b/dimos/control/coordinator.py index c784576c24..2807b2545b 100644 --- a/dimos/control/coordinator.py +++ b/dimos/control/coordinator.py @@ -150,6 +150,8 @@ class ControlCoordinator(Module): # Output: Aggregated joint state for external consumers coordinator_joint_state: Out[JointState] + # Sparse position updates accepted by hardware, for diagnostics and recording. + applied_joint_position_command: Out[JointState] # Input: Streaming joint commands for real-time control joint_command: In[JointState] @@ -903,6 +905,7 @@ def start(self) -> None: task_lock=self._task_lock, joint_to_hardware=self._joint_to_hardware, publish_callback=publish_cb, + publish_command_callback=self.applied_joint_position_command.publish, publish_robot_callback=publish_robot_cb, frame_id=self.config.joint_state_frame_id, log_ticks=self.config.log_ticks, diff --git a/dimos/control/test_control.py b/dimos/control/test_control.py index bff9528cf8..6f40c3969c 100644 --- a/dimos/control/test_control.py +++ b/dimos/control/test_control.py @@ -1132,7 +1132,9 @@ def test_unready_whole_body_is_excluded_from_read_and_write(self, mocker): state, per_hardware = tick_loop._read_all_hardware() imu = tick_loop._read_all_imu() - tick_loop._write_all_hardware({"g1": ({"g1/joint1": 0.25}, ControlMode.SERVO_POSITION)}) + tick_loop._write_all_hardware( + {"g1": ({"g1/joint1": 0.25}, ControlMode.SERVO_POSITION)}, timestamp=123.0 + ) assert state.joint_positions == {} assert per_hardware == {} @@ -1163,7 +1165,9 @@ def test_ready_whole_body_reads_and_writes(self): ) state, _per_hardware = tick_loop._read_all_hardware() - tick_loop._write_all_hardware({"g1": ({"g1/joint1": 0.25}, ControlMode.SERVO_POSITION)}) + tick_loop._write_all_hardware( + {"g1": ({"g1/joint1": 0.25}, ControlMode.SERVO_POSITION)}, timestamp=123.0 + ) assert state.joint_positions == {"g1/joint1": 0.5} adapter.write_motor_commands.assert_called_once() @@ -1296,12 +1300,97 @@ def test_write_all_hardware_rejected_command_logs_error(self, mocker): joint_to_hardware={"arm/joint1": "arm"}, ) - tick_loop._write_all_hardware({"arm": ({"arm/joint1": 0.25}, ControlMode.SERVO_POSITION)}) + tick_loop._write_all_hardware( + {"arm": ({"arm/joint1": 0.25}, ControlMode.SERVO_POSITION)}, timestamp=123.0 + ) log_error.assert_called_once_with( "Hardware arm rejected SERVO_POSITION command from control task" ) + @pytest.mark.parametrize( + "mode,ready,accepted,error,expected", + [ + (ControlMode.POSITION, True, True, None, [0.25]), + (ControlMode.SERVO_POSITION, True, True, None, [0.25]), + (ControlMode.VELOCITY, True, True, None, []), + (ControlMode.TORQUE, True, True, None, []), + (ControlMode.POSITION, False, True, None, []), + (ControlMode.POSITION, True, False, None, []), + (ControlMode.POSITION, True, True, RuntimeError("write failed"), []), + ], + ) + def test_command_feedback_reports_only_successful_position_writes( + self, mocker, mode, ready, accepted, error, expected + ): + hardware = mocker.Mock(spec=ConnectedHardware) + hardware.ready_for_control.return_value = ready + hardware.write_command.return_value = accepted + hardware.write_command.side_effect = error + received = [] + lock = threading.Lock() + + def receive(message): + assert not lock.locked() + received.append(message) + + tick_loop = TickLoop( + tick_rate=100.0, + hardware={"arm": hardware}, + hardware_lock=lock, + tasks={}, + task_lock=threading.Lock(), + joint_to_hardware={}, + publish_command_callback=receive, + frame_id="robot", + ) + tick_loop._write_all_hardware({"arm": ({"arm/joint1": 0.25}, mode)}, timestamp=123.0) + + assert [value for msg in received for value in msg.position] == expected + assert [msg.name for msg in received] == ([["arm/joint1"]] if expected else []) + assert [(msg.ts, msg.frame_id) for msg in received] == ( + [(123.0, "robot")] if expected else [] + ) + assert hardware.write_command.call_count == int(ready) + + def test_command_feedback_combines_successful_writes_and_keeps_sparse_updates(self, mocker): + hardware = { + name: mocker.Mock(spec=ConnectedHardware) for name in ("left", "right", "failed") + } + for hw in hardware.values(): + hw.ready_for_control.return_value = True + hw.write_command.return_value = True + hardware["failed"].write_command.return_value = False + received = [] + tick_loop = TickLoop( + tick_rate=100.0, + hardware=hardware, + hardware_lock=threading.Lock(), + tasks={}, + task_lock=threading.Lock(), + joint_to_hardware={}, + publish_command_callback=received.append, + ) + tick_loop._write_all_hardware( + { + "left": ({"left/joint1": 0.25}, ControlMode.POSITION), + "failed": ({"failed/joint1": 0.75}, ControlMode.POSITION), + "right": ({"right/gripper": 0.5}, ControlMode.SERVO_POSITION), + }, + timestamp=123.0, + ) + tick_loop._write_all_hardware( + { + "right": ({"right/gripper": 0.6}, ControlMode.SERVO_POSITION), + }, + timestamp=124.0, + ) + + assert [(msg.name, msg.position) for msg in received] == [ + (["left/joint1", "right/gripper"], [0.25, 0.5]), + (["right/gripper"], [0.6]), + ] + class TestIntegration: def test_full_trajectory_execution(self, mock_adapter, wait_until): diff --git a/dimos/control/test_per_robot_joint_states.py b/dimos/control/test_per_robot_joint_states.py index bfec8d9f6f..39f8df978a 100644 --- a/dimos/control/test_per_robot_joint_states.py +++ b/dimos/control/test_per_robot_joint_states.py @@ -169,7 +169,10 @@ def test_plain_coordinator_needs_no_per_robot_ports(self, make_coordinator): coordinator.start() assert sorted(coordinator.list_hardware()) == ["left_arm", "right_arm"] - assert sorted(coordinator.outputs) == ["coordinator_joint_state"] + assert sorted(coordinator.outputs) == [ + "applied_joint_position_command", + "coordinator_joint_state", + ] class TestPerRobotPublishing: diff --git a/dimos/control/tick_loop.py b/dimos/control/tick_loop.py index 668186922d..cd191cddab 100644 --- a/dimos/control/tick_loop.py +++ b/dimos/control/tick_loop.py @@ -39,6 +39,7 @@ JointStateSnapshot, ResourceClaim, ) +from dimos.hardware.manipulators.spec import ControlMode from dimos.msgs.sensor_msgs.JointState import JointState from dimos.utils.logging_config import setup_logger @@ -47,7 +48,6 @@ from dimos.control.components import HardwareId, JointName, JointState as JointReading, TaskName from dimos.control.hardware_interface import ConnectedHardware - from dimos.hardware.manipulators.spec import ControlMode from dimos.hardware.whole_body.spec import IMUState logger = setup_logger() @@ -96,6 +96,7 @@ def __init__( task_lock: threading.Lock, joint_to_hardware: dict[JointName, HardwareId], publish_callback: Callable[[JointState], None] | None = None, + publish_command_callback: Callable[[JointState], None] | None = None, publish_robot_callback: Callable[[HardwareId, JointState], None] | None = None, frame_id: str = "coordinator", log_ticks: bool = False, @@ -107,6 +108,7 @@ def __init__( self._task_lock = task_lock self._joint_to_hardware = joint_to_hardware self._publish_callback = publish_callback + self._publish_command_callback = publish_command_callback self._publish_robot_callback = publish_robot_callback self._frame_id = frame_id self._log_ticks = log_ticks @@ -189,7 +191,7 @@ def _tick(self) -> None: hw_commands = self._route_to_hardware(joint_commands) - self._write_all_hardware(hw_commands) + self._write_all_hardware(hw_commands, joint_states.timestamp) if self._publish_callback: self._publish_joint_state(joint_states) @@ -407,9 +409,12 @@ def _route_to_hardware( def _write_all_hardware( self, hw_commands: dict[str, tuple[dict[str, float], ControlMode]], + timestamp: float, ) -> None: - """Write commands to all hardware interfaces.""" + """Write commands and publish accepted position updates outside the hardware lock.""" hardware = self._hardware + names: list[str] = [] + values: list[float] = [] with self._hardware_lock: for hw_id, (positions, mode) in hw_commands.items(): if hw_id in hardware: @@ -421,8 +426,23 @@ def _write_all_hardware( logger.error( f"Hardware {hw_id} rejected {mode.name} command from control task" ) + elif mode in (ControlMode.POSITION, ControlMode.SERVO_POSITION): + names.extend(positions) + values.extend(positions.values()) except Exception as e: logger.error(f"Failed to write to {hw_id}: {e}") + if not names or self._publish_command_callback is None: + return + self._publish_command_callback( + JointState( + ts=timestamp, + frame_id=self._frame_id, + name=names, + position=values, + velocity=[], + effort=[], + ) + ) def _publish_joint_state(self, snapshot: JointStateSnapshot) -> None: """Publish aggregated JointState for external consumers.""" diff --git a/dimos/core/test_build_native_modules.py b/dimos/core/test_build_native_modules.py index 5914df1e8a..a12389e7ed 100644 --- a/dimos/core/test_build_native_modules.py +++ b/dimos/core/test_build_native_modules.py @@ -70,6 +70,7 @@ class _ClassDef(NamedTuple): bases: tuple[str, ...] command: str | None # build_command literal defined in this class body command_kind: str # "absent" | "literal" | "opaque" + owns_cwd: bool = False def _base_names(node: ast.ClassDef) -> tuple[str, ...]: @@ -83,7 +84,7 @@ def _base_names(node: ast.ClassDef) -> tuple[str, ...]: return tuple(names) -def _own_build_command(node: ast.ClassDef) -> tuple[str, str | None]: +def _own_default(node: ast.ClassDef, field: str) -> tuple[str, str | None]: for stmt in node.body: if isinstance(stmt, ast.AnnAssign) and isinstance(stmt.target, ast.Name): target, value = stmt.target.id, stmt.value @@ -95,7 +96,7 @@ def _own_build_command(node: ast.ClassDef) -> tuple[str, str | None]: target, value = stmt.targets[0].id, stmt.value else: continue - if target != "build_command" or value is None: + if target != field or value is None: continue if isinstance(value, ast.Constant) and ( value.value is None or isinstance(value.value, str) @@ -117,15 +118,21 @@ def _scan_all_config_classes() -> list[_ClassDef]: rel = path.relative_to(DIMOS_PROJECT_ROOT).as_posix() for node in ast.walk(ast.parse(path.read_text(), filename=rel)): if isinstance(node, ast.ClassDef): - kind, command = _own_build_command(node) - classes.append(_ClassDef(rel, node.name, _base_names(node), command, kind)) + kind, command = _own_default(node, "build_command") + cwd_kind, _ = _own_default(node, "cwd") + classes.append( + _ClassDef( + rel, node.name, _base_names(node), command, kind, cwd_kind != "absent" + ) + ) return classes def _closure_nix_configs(classes: list[_ClassDef]) -> set[tuple[str, str]]: - """(file, class) for every transitive NativeModuleConfig subclass whose - effective build_command default (own, or inherited from another config in - the closure) mentions nix.""" + """Build owners, including indirect configs that change a command or directory. + + Unmodified inherited builds are covered by their defining ancestor. + """ by_name: dict[str, list[_ClassDef]] = {} for cls in classes: by_name.setdefault(cls.name, []).append(cls) @@ -140,16 +147,18 @@ def _closure_nix_configs(classes: list[_ClassDef]) -> set[tuple[str, str]]: break closure |= added - def effective_command(cls: _ClassDef, seen: frozenset[str]) -> tuple[str, str | None]: + def effective_command( + cls: _ClassDef, seen: frozenset[str] + ) -> tuple[str, str | None, _ClassDef]: if cls.command_kind != "absent": - return cls.command_kind, cls.command + return cls.command_kind, cls.command, cls for base in cls.bases: if base in closure and base != "NativeModuleConfig" and base not in seen: for parent in by_name.get(base, []): - kind, command = effective_command(parent, seen | {cls.name}) + kind, command, owner = effective_command(parent, seen | {cls.name}) if kind != "absent": - return kind, command - return "absent", None + return kind, command, cls if cls.owns_cwd else owner + return "absent", None, cls nix_configs = set() for cls in classes: @@ -164,21 +173,21 @@ def effective_command(cls: _ClassDef, seen: frozenset[str]) -> tuple[str, str | " from EXTERNALLY_PROVISIONED in bin/build-native-modules" ) continue - kind, command = effective_command(cls, frozenset()) + kind, command, owner = effective_command(cls, frozenset()) assert kind != "opaque", ( f"{cls.file}: {cls.name}.build_command must default to a plain string literal " "so bin/build-native-modules can read it without importing dimos" ) - if _SCRIPT.is_nix_build(command): - nix_configs.add((cls.file, cls.name)) + # Deliberately independent of the production command parser: options + # before `build` must not silently remove a config from both scans. + tokens = command.split() if command else [] + if "nix" in tokens and "build" in tokens and "develop" not in tokens: + nix_configs.add((owner.file, owner.name)) return nix_configs -def test_discovery_is_complete_and_flat() -> None: - """The script's direct-base discovery must find every config the transitive - closure finds. A mismatch means a module (e.g. a depth-2 subclass) would - silently escape the publish gate: flatten the hierarchy, or extend the - script's discovery to match.""" +def test_discovery_covers_every_build() -> None: + """Every distinct native build must participate in the publish gate.""" expected = _closure_nix_configs(_scan_all_config_classes()) discovered = { (module.source, module.qualname.rsplit(".", 1)[-1]) for module in _SCRIPT.discover() @@ -187,6 +196,36 @@ def test_discovery_is_complete_and_flat() -> None: assert discovered, "expected at least one nix-built native module" +def test_recorder_is_in_the_publish_manifest() -> None: + recorder = next( + module + for module in _SCRIPT.discover() + if module.qualname == "dimos.experimental.memory.rust_recorder.RustRecorderConfig" + ) + assert recorder.build_dir == "dimos/experimental/memory/rust" + assert _SCRIPT._flake_ref_of(recorder) == ".#dimos-memory-recorder" + + +@pytest.mark.parametrize("override", [None, "build_command", "cwd"]) +def test_inherited_build_coverage_tracks_overrides(override: str | None) -> None: + owner = _ClassDef( + "owner.py", "Owner", ("NativeModuleConfig",), "nix build .#owner", "literal", True + ) + child = _ClassDef( + "child.py", + "Child", + ("Owner",), + "nix build .#child" if override == "build_command" else None, + "literal" if override == "build_command" else "absent", + override == "cwd", + ) + grandchild = _ClassDef("grandchild.py", "Grandchild", ("Child",), None, "absent") + expected = {("owner.py", "Owner")} + if override is not None: + expected.add(("child.py", "Child")) + assert _closure_nix_configs([owner, child, grandchild]) == expected + + def test_ast_extraction_matches_runtime() -> None: """The AST-read defaults must equal what pydantic resolves at runtime — this equivalence is what lets CI discover modules without installing dimos. diff --git a/dimos/experimental/isolated_python/module.py b/dimos/experimental/isolated_python/module.py index 1587ba9cef..723b8004e3 100644 --- a/dimos/experimental/isolated_python/module.py +++ b/dimos/experimental/isolated_python/module.py @@ -50,6 +50,19 @@ def isolated_python_run_command(project: Path, *command: str) -> list[str]: return args +def isolated_python_environment(project: Path) -> dict[str, str]: + """Use a cached project environment without inheriting host interpreter pins.""" + env = dict(os.environ) + # The isolated project picks its own interpreter and venv; host pins + # (e.g. setup-uv exporting UV_PYTHON on CI matrix legs) must not leak in. + env.pop("VIRTUAL_ENV", None) + env.pop("UV_PYTHON", None) + env.pop("UV_PROJECT_ENVIRONMENT", None) + project_key = sha256(str(project).encode()).hexdigest()[:16] + env["UV_PROJECT_ENVIRONMENT"] = str(CACHE_DIR / "isolated-python" / project_key / ".venv") + return env + + class IsolatedPythonModuleConfig(NativeModuleConfig): """Process settings for an isolated Python module.""" @@ -147,14 +160,7 @@ def _new_runtime_name(self) -> str: return self._runtime_name def _runtime_env(self) -> dict[str, str]: - env = dict(os.environ) - # The isolated project picks its own interpreter and venv; host pins - # (e.g. setup-uv exporting UV_PYTHON on CI matrix legs) must not leak in. - env.pop("VIRTUAL_ENV", None) - env.pop("UV_PYTHON", None) - env.pop("UV_PROJECT_ENVIRONMENT", None) - project_key = sha256(str(self.runtime_project).encode()).hexdigest()[:16] - env["UV_PROJECT_ENVIRONMENT"] = str(CACHE_DIR / "isolated-python" / project_key / ".venv") + env = isolated_python_environment(self.runtime_project) env.update(self.config.extra_env) return env diff --git a/dimos/experimental/memory/README.md b/dimos/experimental/memory/README.md index 61779f6fe0..18603ba892 100644 --- a/dimos/experimental/memory/README.md +++ b/dimos/experimental/memory/README.md @@ -11,8 +11,9 @@ The recorder is built as a locked Nix package. Nix supplies Rust, CMake, NASM, SQLite, and the native libraries used by TurboJPEG, so none of those tools or development packages need to be installed on the host. -The Python module builds the package automatically on first use. To build it -ahead of time, run: +The Python module resolves the package through Nix before each launch, so its +wire protocol always matches the Python checkout. Nix reuses the cached package +when the native sources have not changed. To build it ahead of time, run: ```bash cd dimos/experimental/memory/rust diff --git a/dimos/experimental/memory/rust/Cargo.toml b/dimos/experimental/memory/rust/Cargo.toml index c0655dc8a4..e6cb60c064 100644 --- a/dimos/experimental/memory/rust/Cargo.toml +++ b/dimos/experimental/memory/rust/Cargo.toml @@ -10,7 +10,7 @@ anyhow = "1" crossbeam-channel = "0.5" crossbeam-utils = "0.8" dimos-module = { path = "../../../../native/rust/dimos-module" } -lcm-msgs = { git = "https://github.com/dimensionalOS/dimos-lcm.git", branch = "rust-codegen" } +lcm-msgs = { git = "https://github.com/dimensionalOS/dimos-lcm.git", rev = "dd2159513ebfaa7ebc5fc32bf60209cd09aa1ca5" } lz4_flex = "0.14.0" mcap = { version = "0.25.0", default-features = false, features = ["zstd"] } rayon = "1" diff --git a/dimos/experimental/memory/rust/flake.nix b/dimos/experimental/memory/rust/flake.nix index 17ddabe72c..75bd7778c9 100644 --- a/dimos/experimental/memory/rust/flake.nix +++ b/dimos/experimental/memory/rust/flake.nix @@ -46,7 +46,7 @@ cargoLock = { lockFile = ../../../../Cargo.lock; outputHashes = { - "dimos-lcm-0.1.0" = "sha256-GGkx4Mn6NYP6KZecmoRLKGWIih/+y8OgNn12DeXX6n8="; + "dimos-lcm-0.1.0" = "sha256-Z0tKEjNb/VIyfJcYTeZjWhWWYMhQ8wtoe8TjhonPUns="; "pointlio-core-0.1.0" = "sha256-iC7nDbEipfi3cViK7fqKiy2hT9ENGi4Ge7L6Wt1W01Q="; }; }; diff --git a/dimos/experimental/memory/rust/src/decoding.rs b/dimos/experimental/memory/rust/src/decoding.rs index 089de093c9..841fadebab 100644 --- a/dimos/experimental/memory/rust/src/decoding.rs +++ b/dimos/experimental/memory/rust/src/decoding.rs @@ -18,6 +18,7 @@ use lcm_msgs::geometry_msgs::{ PointStamped, PoseStamped, PoseWithCovarianceStamped, TwistStamped, TwistWithCovarianceStamped, WrenchStamped, }; +use lcm_msgs::imitation_msgs::EpisodeStatus; use lcm_msgs::nav_msgs::{OccupancyGrid, Odometry, Path}; use lcm_msgs::sensor_msgs::{ CameraInfo, CompressedImage, Image, Imu, JointState, Joy, PointCloud2, @@ -141,6 +142,11 @@ fn source_timestamp(payload_type: &str, data: &[u8], reception_ts: f64) -> Resul .context("invalid LCM foxglove_msgs.CompressedVideo")?; (message.timestamp.sec, message.timestamp.nanosec) } + "dimos.msgs.imitation_msgs.EpisodeStatus.EpisodeStatus" => { + let message = + EpisodeStatus::decode(data).context("invalid LCM imitation_msgs.EpisodeStatus")?; + (message.header.stamp.sec, message.header.stamp.nsec) + } "dimos.msgs.geometry_msgs.Transform.Transform" => { let message = TFMessage::decode(data).context("invalid LCM TFMessage")?; let Some(transform) = message.transforms.first() else { @@ -160,3 +166,36 @@ pub(crate) fn header_timestamp(sec: i32, nsec: i32, fallback: f64) -> f64 { fallback } } + +#[cfg(test)] +mod tests { + use lcm_msgs::imitation_msgs::EpisodeStatus; + use lcm_msgs::std_msgs::{Header, Time}; + + use super::source_timestamp; + + #[test] + fn episode_status_uses_source_timestamp() { + let status = EpisodeStatus { + header: Header { + stamp: Time { + sec: 42, + nsec: 250_000_000, + }, + ..Header::default() + }, + state: "recording".to_string(), + last_event: "start".to_string(), + ..EpisodeStatus::default() + }; + + let timestamp = source_timestamp( + "dimos.msgs.imitation_msgs.EpisodeStatus.EpisodeStatus", + &status.encode(), + 99.0, + ) + .expect("generated status should decode"); + + assert_eq!(timestamp, 42.25); + } +} diff --git a/dimos/experimental/memory/rust_recorder.py b/dimos/experimental/memory/rust_recorder.py index 667945cce6..5cc5cc20cb 100644 --- a/dimos/experimental/memory/rust_recorder.py +++ b/dimos/experimental/memory/rust_recorder.py @@ -20,7 +20,7 @@ from pathlib import Path from typing import Annotated, Any, Literal, TypeAlias -from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator +from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, field_validator, model_validator from dimos.constants import DIMOS_PROJECT_ROOT from dimos.core.core import rpc @@ -86,18 +86,13 @@ class RustMcapStoreConfig(RustStoreConfig): class RustRecorderConfig(NativeModuleConfig): - """Compatibility-first configuration for :class:`RustRecorder`. + """Record connected streams to a Rust-backed SQLite or MCAP artifact. Python owns artifact lifecycle and stream registration. The native process receives only ``store``, ``encoding_threads``, and the internally resolved ``streams`` list over stdin. """ - executable: str = "result/bin/dimos-memory-recorder" - build_command: str = "nix build -L .#dimos-memory-recorder" - cwd: str = "rust" - stdin_config: bool = True - store: RustRecordingStoreConfig = Field( default_factory=RustSqliteStoreConfig, description="Record-only native storage backend and artifact path.", @@ -113,6 +108,12 @@ class RustRecorderConfig(NativeModuleConfig): exclude=True, description="Maximum rotated artifacts retained when on_existing is backup.", ) + + executable: str = "result/bin/dimos-memory-recorder" + build_command: str = "nix build -L .#dimos-memory-recorder" + cwd: str = "rust" + stdin_config: bool = True + record_tf: bool = Field( default=True, exclude=True, @@ -133,11 +134,14 @@ class RustRecorderConfig(NativeModuleConfig): ge=1, description="CPU workers for transport decoding and storage encoding.", ) - streams: list[RustStreamSpec] = Field( - default_factory=list, - init=False, - description="Resolved stream plan populated internally before native launch.", - ) + _streams: list[RustStreamSpec] = PrivateAttr(default_factory=list) + + @model_validator(mode="after") + def _resolve_cwd(self) -> RustRecorderConfig: + # Subclassed recorders share this native project, regardless of their source file. + if not Path(self.cwd).is_absolute(): + self.cwd = str(Path(__file__).parent / self.cwd) + return self @model_validator(mode="after") def _stdin_only(self) -> RustRecorderConfig: @@ -145,6 +149,13 @@ def _stdin_only(self) -> RustRecorderConfig: raise ValueError("RustRecorder is stdin-only and does not accept extra_args") return self + def to_config_dict(self) -> dict[str, Any]: + return { + "store": self.store.model_dump(), + "encoding_threads": self.encoding_threads, + "streams": [stream.model_dump() for stream in self._streams], + } + class RustRecorder(NativeModule): """Experimentally record connected ``In`` ports to native SQLite or MCAP. @@ -185,7 +196,7 @@ def start(self) -> None: return self._prepare_store(specs) - self.config.streams = specs + self.config._streams = specs super().start() def _stream_specs(self) -> list[RustStreamSpec]: @@ -225,6 +236,11 @@ def _stream_specs(self) -> list[RustStreamSpec]: logger.warning("Native recorder has no connected streams") return specs + def _collect_topics(self) -> dict[str, str]: + topics = super()._collect_topics() + enabled_ports = {spec.port for spec in self.config._streams} + return {port: topic for port, topic in topics.items() if port in enabled_ports} + @staticmethod def _default_codec(payload_type: type[Any]) -> str: if issubclass(payload_type, Image): @@ -271,13 +287,12 @@ def _prepare_store(self, specs: list[RustStreamSpec]) -> None: path.parent.mkdir(parents=True, exist_ok=True) if self.config.store.kind == "mcap": return - - with SqliteStore(path=str(path)) as store: - if self.config.on_existing is OnExisting.APPEND: + if self.config.on_existing is OnExisting.APPEND: + with SqliteStore(path=str(path)) as store: existing = set(store.list_streams()) for name in {spec.name for spec in specs}.intersection(existing): store.delete_stream(name) - + with SqliteStore(path=str(path)) as store: ports = self.inputs for spec in specs: tf_payload_type = f"{TFMessage.__module__}.{TFMessage.__qualname__}" diff --git a/dimos/experimental/memory/test_rust_recorder.py b/dimos/experimental/memory/test_rust_recorder.py index c8a8e40997..3e4c6bbf32 100644 --- a/dimos/experimental/memory/test_rust_recorder.py +++ b/dimos/experimental/memory/test_rust_recorder.py @@ -89,7 +89,7 @@ def test_specs_use_native_defaults_remapping_and_configured_workers( connect(recorder, color_image="/camera", odometry="/odom") specs = recorder._stream_specs() - recorder.config.streams = specs + recorder.config._streams = specs config = recorder.config.to_config_dict() assert config["encoding_threads"] == 7 @@ -186,8 +186,8 @@ def test_default_store_path_is_resolved_from_the_project_root() -> None: def test_native_recorder_is_built_and_run_from_the_nix_package() -> None: config = RustRecorderConfig() - assert config.cwd == "rust" - assert config.build_command == "nix build -L .#dimos-memory-recorder" + assert Path(config.cwd) == Path(__file__).with_name("rust") + assert config.build_command == ("nix build -L .#dimos-memory-recorder") assert config.executable == "result/bin/dimos-memory-recorder" @@ -245,7 +245,7 @@ def test_mcap_store_uses_python_codec_defaults_and_does_not_precreate_the_artifa assert [spec.codec for spec in specs] == ["jpeg", "lcm"] assert not path.exists() - recorder.config.streams = specs + recorder.config._streams = specs assert recorder.config.to_config_dict()["store"] == { "kind": "mcap", "path": str(path), diff --git a/dimos/experimental/memory/test_rust_recorder_e2e.py b/dimos/experimental/memory/test_rust_recorder_e2e.py index d26a261651..a8b187ab4e 100644 --- a/dimos/experimental/memory/test_rust_recorder_e2e.py +++ b/dimos/experimental/memory/test_rust_recorder_e2e.py @@ -40,6 +40,10 @@ RustRecordingStoreConfig, RustSqliteStoreConfig, ) +from dimos.imitation.collection.profile import CollectionFeature, CollectionProfile +from dimos.imitation.collection.recorder import collection_recorder +from dimos.imitation.collection.recording import RecordingSchema +from dimos.imitation.dataprep.core import SyncConfig from dimos.memory.codecs.lcm import LcmCodec from dimos.memory.codecs.lz4 import Lz4Codec from dimos.memory.store.mcap import McapStore @@ -140,6 +144,55 @@ def test_rust_artifact_is_readable_by_python_memory2( store_kind: str, monkeypatch: pytest.MonkeyPatch, ) -> None: + _capture_native_artifact(tmp_path, rust_recorder_executable, store_kind, monkeypatch) + + +@pytest.mark.parametrize("store_kind", ["sqlite", "mcap"]) +def test_native_collection_directory_preserves_schema_and_messages( + tmp_path, + rust_recorder_executable, + store_kind, + monkeypatch, +): + camera = CollectionFeature( + stream="color_image", + message_type=Image, + field="data", + dtype="video", + shape=(16, 16, 3), + names=["height", "width", "channels"], + ) + state = CollectionFeature( + stream="imu", + message_type=Imu, + field="angular_velocity", + dtype="float32", + shape=(3,), + names=["x", "y", "z"], + ) + profile = CollectionProfile( + name="native-test", + robot_type="test", + observations={"image": camera, "state": state}, + actions={"action": state}, + sync=SyncConfig(anchor="image", rate_hz=30, tolerance_ms=20), + ) + artifact = _capture_native_artifact( + tmp_path, rust_recorder_executable, store_kind, monkeypatch, profile + ) + schema = RecordingSchema.read(artifact.parent) + assert schema.payload == artifact.name + assert schema.observation["state"].stream == "imu" + assert schema.name == "native-test" + + +def _capture_native_artifact( + tmp_path: Path, + rust_recorder_executable: Path, + store_kind: str, + monkeypatch: pytest.MonkeyPatch, + profile: CollectionProfile | None = None, +) -> Path: suffix = ".db" if store_kind == "sqlite" else ".mcap" artifact = tmp_path / f"recording{suffix}" store: RustRecordingStoreConfig @@ -165,6 +218,20 @@ def test_rust_artifact_is_readable_by_python_memory2( connect_timeout=0, ), ) + if profile is not None: + kwargs = { + "executable": str(rust_recorder_executable), + "encoding_threads": 2, + "stream_codecs": {"imu": "lz4+lcm"}, + "session": recorder.config.session, + } + recorder.stop() + atom = collection_recorder( + profile=profile, recording=tmp_path / "session", format=store_kind + ).active_blueprints[0] + recorder = atom.module(**atom.kwargs, **kwargs) + artifact = Path(recorder.config.store.path) + recorder.status.transport = FakeTransport("dimos/collection_status") session_pool = ZenohSessionPool() channel_suffix = uuid.uuid4().hex[:8] publisher: ZenohTransport[Imu] = ZenohTransport( @@ -191,8 +258,8 @@ def test_rust_artifact_is_readable_by_python_memory2( recorder.color_image.transport = FakeTransport(image_topic) # type: ignore[assignment] specs = recorder._stream_specs() recorder._prepare_store(specs) - recorder.config.streams = specs - launch = recorder._stdin_blob({"imu": imu_topic, "color_image": image_topic}) + recorder.config._streams = specs + launch = recorder._stdin_blob(recorder._collect_topics()) env = { **os.environ, @@ -261,6 +328,8 @@ def test_rust_artifact_is_readable_by_python_memory2( < 5 ) + return artifact + @pytest.mark.parametrize( "store_kind", @@ -351,7 +420,7 @@ def test_tf_records_over_zenoh_and_replays_through_python( recorder.tf.transport = FakeTransport(publisher.channel) # type: ignore[assignment] specs = recorder._stream_specs() recorder._prepare_store(specs) - recorder.config.streams = specs + recorder.config._streams = specs launch = recorder._stdin_blob({"tf": publisher.channel}) env = { diff --git a/dimos/hardware/sensors/camera/module.py b/dimos/hardware/sensors/camera/module.py index 83020165da..75892a99ed 100644 --- a/dimos/hardware/sensors/camera/module.py +++ b/dimos/hardware/sensors/camera/module.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from collections.abc import Callable import time from pydantic import Field @@ -24,7 +23,7 @@ from dimos.core.module import Module, ModuleConfig from dimos.core.stream import Out from dimos.hardware.sensors.camera.spec import CameraHardware -from dimos.hardware.sensors.camera.webcam import Webcam +from dimos.hardware.sensors.camera.webcam import Webcam, WebcamConfig from dimos.msgs.geometry_msgs.Quaternion import Quaternion from dimos.msgs.geometry_msgs.Transform import Transform from dimos.msgs.geometry_msgs.Vector3 import Vector3 @@ -47,7 +46,7 @@ def default_transform() -> Transform: class CameraModuleConfig(ModuleConfig): frame_id: str = "camera_link" transform: Transform | None = Field(default_factory=default_transform) - hardware: Callable[[], CameraHardware] | CameraHardware = Webcam + hardware: WebcamConfig = Field(default_factory=WebcamConfig) frequency: float = 0.0 # Hz, 0 means no limit @@ -63,10 +62,7 @@ class CameraModule(Module, perception.Camera): def start(self) -> None: super().start() - if callable(self.config.hardware): - self.hardware = self.config.hardware() - else: - self.hardware = self.config.hardware + self.hardware = Webcam(**self.config.hardware.model_dump()) stream = self.hardware.image_stream() @@ -103,8 +99,9 @@ def publish_metadata(self) -> None: @rpc def stop(self) -> None: - if self.hardware and hasattr(self.hardware, "stop"): - self.hardware.stop() + hardware = getattr(self, "hardware", None) + if hardware is not None and hasattr(hardware, "stop"): + hardware.stop() super().stop() diff --git a/dimos/hardware/sensors/camera/realsense/rust/Cargo.lock b/dimos/hardware/sensors/camera/realsense/rust/Cargo.lock index ad82645d0d..a6152609bd 100644 --- a/dimos/hardware/sensors/camera/realsense/rust/Cargo.lock +++ b/dimos/hardware/sensors/camera/realsense/rust/Cargo.lock @@ -542,7 +542,7 @@ dependencies = [ [[package]] name = "dimos-lcm" version = "0.1.0" -source = "git+https://github.com/dimensionalOS/dimos-lcm.git?branch=rust-codegen#04d78e8622500244123ba9cefa4c51b4cb454549" +source = "git+https://github.com/dimensionalOS/dimos-lcm.git?rev=dd2159513ebfaa7ebc5fc32bf60209cd09aa1ca5#dd2159513ebfaa7ebc5fc32bf60209cd09aa1ca5" dependencies = [ "byteorder", "socket2 0.5.10", @@ -555,7 +555,7 @@ version = "0.1.0" dependencies = [ "dimos-lcm", "dimos-module-macros", - "lcm-msgs", + "lcm-msgs 0.1.0 (git+https://github.com/dimensionalOS/dimos-lcm.git?rev=dd2159513ebfaa7ebc5fc32bf60209cd09aa1ca5)", "nalgebra", "rayon", "serde", @@ -583,7 +583,7 @@ name = "dimos-realsense" version = "0.1.0" dependencies = [ "dimos-module", - "lcm-msgs", + "lcm-msgs 0.1.0 (git+https://github.com/dimensionalOS/dimos-lcm.git?branch=rust-codegen)", "pkg-config", "realsense-rust", "realsense-sys", @@ -1363,6 +1363,14 @@ dependencies = [ "byteorder", ] +[[package]] +name = "lcm-msgs" +version = "0.1.0" +source = "git+https://github.com/dimensionalOS/dimos-lcm.git?rev=dd2159513ebfaa7ebc5fc32bf60209cd09aa1ca5#dd2159513ebfaa7ebc5fc32bf60209cd09aa1ca5" +dependencies = [ + "byteorder", +] + [[package]] name = "libc" version = "0.2.189" diff --git a/dimos/hardware/sensors/camera/test_webcam.py b/dimos/hardware/sensors/camera/test_webcam.py index 6e390d078c..9a4727cd7e 100644 --- a/dimos/hardware/sensors/camera/test_webcam.py +++ b/dimos/hardware/sensors/camera/test_webcam.py @@ -12,10 +12,83 @@ # See the License for the specific language governing permissions and # limitations under the License. -from dimos.hardware.sensors.camera.webcam import Webcam +from unittest.mock import call + +import cv2 +import numpy as np +import pytest_mock + +from dimos.hardware.sensors.camera.module import CameraModule, CameraModuleConfig +from dimos.hardware.sensors.camera.webcam import Webcam, WebcamConfig from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo +def test_camera_config_defaults_to_one_typed_webcam_field() -> None: + config = CameraModuleConfig() + + assert isinstance(config.hardware, WebcamConfig) + + +def test_camera_config_parses_nested_webcam_settings() -> None: + config = CameraModuleConfig.model_validate({"hardware": {"camera_index": 2}}) + + assert isinstance(config.hardware, WebcamConfig) + assert config.hardware.camera_index == 2 + + +def test_camera_can_stop_before_start() -> None: + camera = CameraModule() + + camera.stop() + + assert camera._module_closed + + +def test_start_requests_and_reports_camera_mode(mocker: pytest_mock.MockerFixture) -> None: + capture = mocker.patch("cv2.VideoCapture").return_value + capture.isOpened.return_value = True + capture.get.side_effect = [30.0, 640.0, 480.0] + thread = mocker.patch("dimos.hardware.sensors.camera.webcam.threading.Thread").return_value + log = mocker.patch("dimos.hardware.sensors.camera.webcam.logger.info") + webcam = Webcam(camera_index=2, width=640, height=480, fps=30.0) + + webcam.start() + + assert capture.set.call_args_list == [ + call(cv2.CAP_PROP_FRAME_WIDTH, 640), + call(cv2.CAP_PROP_FRAME_HEIGHT, 480), + call(cv2.CAP_PROP_FPS, 30.0), + ] + log.assert_called_once_with( + "Webcam %s requested %.1fHz %dx%d; negotiated %.1fHz %dx%d", + 2, + 30.0, + 640, + 480, + 30.0, + 640, + 480, + ) + thread.start.assert_called_once_with() + + +def test_capture_timestamp_is_taken_immediately_after_read( + mocker: pytest_mock.MockerFixture, +) -> None: + frame = np.zeros((2, 3, 3), dtype=np.uint8) + capture = mocker.MagicMock() + capture.read.return_value = (True, frame) + webcam = Webcam(camera_index=2, width=3, height=2, fps=30.0) + webcam._capture = capture + now = mocker.patch("dimos.hardware.sensors.camera.webcam.time.time", return_value=42.5) + + image = webcam.capture_frame() + + now.assert_called_once_with() + assert image.ts == 42.5 + assert image.data.shape == (2, 3, 3) + + def test_webcam_without_intrinsics_reports_a_nominal_pinhole() -> None: info = Webcam(width=1280, height=720).camera_info assert (info.width, info.height) == (1280, 720) diff --git a/dimos/hardware/sensors/camera/webcam.py b/dimos/hardware/sensors/camera/webcam.py index 8523f356b6..e4eee3dfe3 100644 --- a/dimos/hardware/sensors/camera/webcam.py +++ b/dimos/hardware/sensors/camera/webcam.py @@ -25,8 +25,11 @@ from dimos.hardware.sensors.camera.spec import CameraConfig, CameraHardware from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.utils.logging_config import setup_logger from dimos.utils.reactive import backpressure +logger = setup_logger() + def _parse_camera_device(value: Any) -> Any: if isinstance(value, str): @@ -104,6 +107,18 @@ def start(self): # type: ignore[no-untyped-def] # Set camera properties self._capture.set(cv2.CAP_PROP_FRAME_WIDTH, self.config.width) # type: ignore[attr-defined] self._capture.set(cv2.CAP_PROP_FRAME_HEIGHT, self.config.height) # type: ignore[attr-defined] + if self.config.fps > 0: + self._capture.set(cv2.CAP_PROP_FPS, self.config.fps) # type: ignore[attr-defined] + logger.info( + "Webcam %s requested %.1fHz %dx%d; negotiated %.1fHz %dx%d", + self.config.camera_index, + self.config.fps, + self.config.width, + self.config.height, + self._capture.get(cv2.CAP_PROP_FPS), # type: ignore[attr-defined] + round(self._capture.get(cv2.CAP_PROP_FRAME_WIDTH)), # type: ignore[attr-defined] + round(self._capture.get(cv2.CAP_PROP_FRAME_HEIGHT)), # type: ignore[attr-defined] + ) # Clear stop event and start the capture thread self._stop_event.clear() @@ -136,6 +151,7 @@ def capture_frame(self) -> Image: import cv2 ret, frame = self._capture.read() # type: ignore[attr-defined] + acquired_at = time.time() if not ret: raise RuntimeError(f"Failed to read frame from camera {self.config.camera_index}") @@ -149,7 +165,7 @@ def capture_frame(self) -> Image: frame_rgb, format=ImageFormat.RGB, # We converted to RGB above frame_id=self._frame("camera_optical"), # Standard frame ID for camera images - ts=time.time(), # Current timestamp + ts=acquired_at, ) if self.config.stereo_slice in ("left", "right"): diff --git a/dimos/hardware/sensors/lidar/livox/rust/Cargo.toml b/dimos/hardware/sensors/lidar/livox/rust/Cargo.toml index 5a373b33d9..c99cecdb0c 100644 --- a/dimos/hardware/sensors/lidar/livox/rust/Cargo.toml +++ b/dimos/hardware/sensors/lidar/livox/rust/Cargo.toml @@ -24,7 +24,7 @@ imu = "sensor_msgs.Imu" crc = "3" dimos-module = { path = "../../../../../../native/rust/dimos-module" } etherparse = "0.21.0" -lcm-msgs = { git = "https://github.com/dimensionalOS/dimos-lcm.git", branch = "rust-codegen" } +lcm-msgs = { git = "https://github.com/dimensionalOS/dimos-lcm.git", rev = "dd2159513ebfaa7ebc5fc32bf60209cd09aa1ca5" } pcap-parser = "0.17.0" serde = { version = "1", features = ["derive"] } serde_json = "1" diff --git a/dimos/hardware/sensors/lidar/pointlio/rust/Cargo.toml b/dimos/hardware/sensors/lidar/pointlio/rust/Cargo.toml index 6a1af9d487..90d31e9117 100644 --- a/dimos/hardware/sensors/lidar/pointlio/rust/Cargo.toml +++ b/dimos/hardware/sensors/lidar/pointlio/rust/Cargo.toml @@ -33,7 +33,7 @@ odometry = "nav_msgs.Odometry" dimos-livox = { path = "../../livox/rust" } pointlio-core = { git = "https://github.com/dimensionalOS/dimos-module-pointlio.git", rev = "3c1dcf8b9dfe6402aaf141e116f4caef6c5c8088" } dimos-module = { path = "../../../../../../native/rust/dimos-module" } -lcm-msgs = { git = "https://github.com/dimensionalOS/dimos-lcm.git", branch = "rust-codegen" } +lcm-msgs = { git = "https://github.com/dimensionalOS/dimos-lcm.git", rev = "dd2159513ebfaa7ebc5fc32bf60209cd09aa1ca5" } serde = { version = "1", features = ["derive"] } serde_json = "1" tokio = { version = "1", features = ["rt-multi-thread", "macros"] } diff --git a/dimos/imitation/README.md b/dimos/imitation/README.md index 33901320a4..46f764dbbf 100644 --- a/dimos/imitation/README.md +++ b/dimos/imitation/README.md @@ -1,126 +1,79 @@ -# Imitation Learning +# Imitation learning -Collect demonstrations, build training datasets, and run trained policies in -DimOS. Teleoperation records episodes to a session DB, and DataPrep converts -that DB into a LeRobot or HDF5 dataset for imitation learning. +Collection uses ordinary DimOS Blueprints. The graph owns robot hardware, +cameras, transports, and runtime lifecycle. A `CollectionProfile` declares +typed inputs and dataset projections; `collection_recorder(profile=...)` +creates matching recorder ports before autoconnect. Import the factory from +`dimos.imitation.collection.recorder`. `CollectionRecorder` extends `RustRecorder` +with collection directory and schema preparation; both use the same Rust executable. -``` -teleop (WebXR) ─▶ CollectionRecorder ─▶ session__.db ─▶ dimos dataprep ─▶ dataset -``` - -After training, use the production -[`LeRobotPolicyModule`](policy/lerobot/README.md) to run a checkpoint against -live camera and joint-state observations. - ---- - -## 1. Record a session - -Run a collection blueprint. Add `--simulation` to drive MuJoCo; omit it for real -hardware (a RealSense + the arm). - -```bash -# XArm7 in sim -dimos --simulation run learning-collect-webxr-xarm7 - -# Piper on real hardware -dimos run learning-collect-webxr-piper -``` - -This brings up teleop, a RealSense (real only), the episode monitor, and the -recorder, all wired together. - -### Controls (WebXR) - -| Button | Action | -| --- | --- | -| **A** (right) / **X** (left) | **Hold to engage** — the arm tracks the controller only while held | -| **B** | **Toggle record** — press to start an episode, press again to save it | -| **Y** | **Discard** the in-progress episode | - -So a take is: hold **A** to move the arm into place → press **B** to start → -perform the task → press **B** to save (or **Y** to throw it away). The terminal -prints one line per transition: - -``` -[collect] ▶ RECORDING episode (state=recording saved=0 discarded=0) -[collect] ✓ SAVED episode (state=idle saved=1 discarded=0) -``` - -> End each good take with **B** before quitting — an episode still recording at -> shutdown is dropped. - -### Where the recording goes - -``` -~/.local/state/dimos/recordings/session__.db -``` - -A new timestamped file per run (nothing is overwritten). It records three -streams: `color_image`, `coordinator_joint_state`, and `status` (the episode -start/save/discard markers). - -The exact path is printed when the recorder starts — note it for the next step. +Profiles have no separate registry. `dimos run` discovers Blueprints through the +built-in registry or installed `dimos.blueprints` entry points. The Blueprint +passes a profile to its recorder; the profile name is recording metadata, not +a Blueprint lookup key. Profile validation checks declarations and shared-source +consistency. Recorder wiring checks required inputs; preparation validates the +actual recorded values. ---- - -## 2. Build a dataset - -DataPrep is an offline batch step that reads a session DB and writes a dataset. -The obs/action stream mapping is nested, so it comes from a JSON config — start -from [`dataprep/example_config.json`](dataprep/example_config.json) and edit the -`source`/`output` to taste. +## OpenYAM Quest collection ```bash -# LeRobot v3.0 (default) -dimos dataprep build \ - --source ~/.local/state/dimos/recordings/session_xarm7_20260622_120000.db \ - --config dimos/imitation/dataprep/example_config.json - -# HDF5 instead -dimos dataprep build -s -c -f hdf5 +dimos run openyam-quest-collection \ + --recorder.recording recordings/session-001 \ + --episodes.task "pick up the cube" ``` -`--source` / `--output` / `--format` override whatever the config specifies, so -you can reuse one config across runs and just swap `--source`. The dataset is -written to the config's `output.path` (the example uses `data/datasets/session`) -unless you pass `--output`. - -Inspect the result (features, shapes, dtypes, episode/frame counts): - -```bash -dimos dataprep inspect data/datasets/session # LeRobot dir -dimos dataprep inspect data/datasets/session.hdf5 # HDF5 file +Configure camera and hardware options through `dimos run BLUEPRINT --help`. +Quest B starts/saves an episode; Y discards it. Python clients can use +`Dimos.connect().find_module_by_spec(EpisodeControlSpec)` and its +`get_status()` and `command(event)` RPCs instead. + +The recording is a new directory containing `schema.json` and +`recording.mcap`, or `recording.db` with `--recorder.format sqlite`. +Existing directories are rejected. Copy or move the complete directory. +The inherited `store` settings must match the destination derived from `recording` +and `format`; collection requires `on_existing=error` and does not rotate backups. +The xArm and Piper collection blueprints also use this recorder, with timestamped +session directories and SQLite payloads. +Stopping the runtime leaves an active episode incomplete; export excludes +incomplete and discarded episodes. Support the arms before shutdown. + +## Prepare a recording + +```python +from pathlib import Path + +from dimos.imitation.collection.recording import RecordingSchema +from dimos.imitation.dataprep.core import OutputConfig +from dimos.imitation.dataprep.lerobot import run_lerobot_dataprep + +directory = Path("recordings/session-001") +config = RecordingSchema.read(directory).dataprep_config( + directory, OutputConfig(format="lerobot", path=Path("datasets/session-001")) +) +run_lerobot_dataprep(config) ``` -Each dataset gets a `dimos_meta.json` sidecar recording exactly how it was built -(source, sync, episodes). - ---- - -## 3. Config reference +Preparation uses the saved schema, not a current robot profile. Only prepare +trusted recordings; custom message classes must be installed in the reader +environment. Generic `run_dataprep(config)` supports HDF5 output. -See [`dataprep/example_config.json`](dataprep/example_config.json) for a full, -working example. The fields that matter: +Each feature declares its recorded source's meaning with `source_kind`: -- **`source`** — the session `.db`. -- **`observation` / `action`** — map a dataset feature name to a recorded - `{stream, field}`. Action defaults to the *next* frame's joint state (see - `action_shift`), giving a next-state behavioral-cloning target. -- **`sync`** — resample everything onto one timeline: `anchor` stream, - `rate_hz`, nearest-match `tolerance_ms`, and `action_shift` (1 = next-state BC, - 0 = action == state). `fps` is derived from `rate_hz` unless set explicitly. -- **`output`** — `format` (`lerobot` | `hdf5`), `path`, and `metadata` - (`robot`, `default_task_label`, …). +- `"snapshot"` (default): align to the nearest observation within the configured + tolerance. This also applies when measured state supplies a teaching action. +- `"joint_position_updates"`: reconstruct persistent `JointState.position` + targets by joint name, using only updates at or before each dataset timestamp. + Omitted joints retain their targets, including across episode boundaries. + Missing initial joints and malformed updates fail validation. ---- +Features sharing a recorded stream must declare the same source kind. Command +history is reconstructed once before projecting individual features. Inspection +and preparation share alignment and value checks; MCAP and SQLite capture remain +unaligned, native-rate streams. Start a new recording after an unrecorded target +reset or control-mode change. -## Notes +## Policy execution -- **Sim vs real camera** — under `--simulation` the MuJoCo camera supplies - `color_image`; on real hardware a RealSense does. The blueprint picks the - right one automatically. -- **"action" is the measured next joint state**, not a recorded command. For - true commanded actions you'd record `joint_command` and map `action` to it. -- **Old vs new sessions** — recordings made before the `coordinator_joint_state` - rename use the old stream name; point a matching config at them, or re-record. +The [LeRobot module](policy/lerobot/README.md) provides isolated checkpoint +loading, preflight, and controlled trajectory execution. Collection profiles +do not define arbitrary policy-backend compatibility. diff --git a/dimos/imitation/collection/blueprint.py b/dimos/imitation/collection/blueprint.py index f07d8d627a..5b8ee5a4a7 100644 --- a/dimos/imitation/collection/blueprint.py +++ b/dimos/imitation/collection/blueprint.py @@ -12,34 +12,33 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Recording blueprints. - -`CollectionRecorder` (a memory Recorder) captures the obs/action/status -streams to a SQLite session DB during the run and flushes it durably on -shutdown. DataPrep reads that DB afterwards. -""" +"""Profile-based Rust recording for the xArm and Piper teleop blueprints.""" from __future__ import annotations from datetime import datetime +from pathlib import Path +from typing import cast from dimos.constants import RECORDINGS_DIR +from dimos.control.coordinator import ControlCoordinator +from dimos.core.coordination.blueprint_config.fields import module_config_cls from dimos.core.coordination.blueprints import Blueprint, autoconnect from dimos.core.global_config import global_config from dimos.hardware.sensors.camera.realsense.camera import RealSenseCamera +from dimos.hardware.sensors.camera.spec import CameraConfig from dimos.imitation.collection.episode_monitor import EpisodeMonitorModule -from dimos.imitation.collection.recorder import CollectionRecorder +from dimos.imitation.collection.profile import CollectionFeature, CollectionProfile +from dimos.imitation.collection.recorder import collection_recorder +from dimos.imitation.dataprep.core import SyncConfig +from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.sensor_msgs.JointState import JointState from dimos.teleop.webxr.blueprints import ( teleop_webxr_piper, teleop_webxr_xarm7, ) -def _session_db(robot: str) -> str: - """Timestamped session DB path under RECORDINGS_DIR, namespaced by robot.""" - return str(RECORDINGS_DIR / f"session_{robot}_{datetime.now():%Y%m%d_%H%M%S}.db") - - def _camera_if_real() -> tuple[Blueprint, ...]: """Real RealSense only off-sim. In `--simulation` the teleop coordinator's MujocoSimModule already publishes color_image on /camera/color_image, so a @@ -49,30 +48,52 @@ def _camera_if_real() -> tuple[Blueprint, ...]: return (RealSenseCamera.blueprint(enable_pointcloud=False),) -# buttons / color_image / coordinator_joint_state / status are left to -# autoconnect — each name is unique across the composed blueprint, so it -# resolves to a stable / topic shared by producer and recorder. The -# recorder captures whatever joints are present, so the coordinator's aggregate -# stream is its intended input (see dimos/control/README.md). -learning_collect_webxr_xarm7 = autoconnect( - CollectionRecorder.blueprint( - db_path=_session_db("xarm7"), - poseless_streams=["color_image", "coordinator_joint_state", "status"], - record_tf=False, - ), - EpisodeMonitorModule.blueprint(), # default button_map: toggle=B, discard=Y - teleop_webxr_xarm7, - *_camera_if_real(), -) +def _collection_components(robot: str, teleop: Blueprint) -> tuple[Blueprint, ...]: + producers = autoconnect(teleop, *_camera_if_real()) + coordinator = next( + atom for atom in producers.active_blueprints if issubclass(atom.module, ControlCoordinator) + ) + joints = [joint for hardware in coordinator.kwargs["hardware"] for joint in hardware.joints] + camera = next( + atom + for atom in producers.active_blueprints + if any( + stream.name == "color_image" and stream.direction == "out" for stream in atom.streams + ) + ) + camera_config = cast("CameraConfig", module_config_cls(camera)(**camera.kwargs)) + state = CollectionFeature( + stream="coordinator_joint_state", + message_type=JointState, + field="position", + dtype="float32", + shape=(len(joints),), + names=joints, + ) + profile = CollectionProfile( + name=f"{robot}-webxr", + robot_type=robot, + observations={ + "camera": CollectionFeature( + stream="color_image", + message_type=Image, + field="data", + dtype="video", + shape=(camera_config.height, camera_config.width, 3), + names=["height", "width", "channels"], + ), + "state": state, + }, + actions={"action": state.model_copy(deep=True)}, + sync=SyncConfig(anchor="camera", rate_hz=camera_config.fps, tolerance_ms=50), + ) + directory: Path = RECORDINGS_DIR / f"session_{robot}_{datetime.now():%Y%m%d_%H%M%S_%f}" + return ( + collection_recorder(profile=profile, recording=directory, format="sqlite"), + EpisodeMonitorModule.blueprint(), + producers, + ) -learning_collect_webxr_piper = autoconnect( - CollectionRecorder.blueprint( - db_path=_session_db("piper"), - poseless_streams=["color_image", "coordinator_joint_state", "status"], - record_tf=False, - ), - EpisodeMonitorModule.blueprint(), # default button_map: toggle=B, discard=Y - teleop_webxr_piper, - *_camera_if_real(), -) +learning_collect_webxr_xarm7 = autoconnect(*_collection_components("xarm7", teleop_webxr_xarm7)) +learning_collect_webxr_piper = autoconnect(*_collection_components("piper", teleop_webxr_piper)) diff --git a/dimos/imitation/collection/episode_monitor.py b/dimos/imitation/collection/episode_monitor.py index 03e171bbfc..f00516e862 100644 --- a/dimos/imitation/collection/episode_monitor.py +++ b/dimos/imitation/collection/episode_monitor.py @@ -12,28 +12,33 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Single point of teleop-input → EpisodeStatus translation. +"""Single point of operator-input → EpisodeStatus translation. -Watches buttons / keyboard, runs the start/save/discard state machine, +Watches Quest buttons and accepts RPC commands, runs the episode state machine, publishes EpisodeStatus on every transition. RecordReplay (or whatever -records the bus) captures that stream into session.db; DataPrep reads -only the recorded EpisodeStatus events offline — never raw buttons or -keypresses. +records the bus) captures that stream into session.db; DataPrep reads only +the recorded EpisodeStatus events offline — never raw operator input. """ from __future__ import annotations import threading import time -from typing import Any, Literal, TypeAlias +from typing import Any, Literal, Protocol, TypeAlias -from pydantic import BaseModel, Field, field_validator +from pydantic import Field, field_validator from reactivex.abc import DisposableBase from reactivex.disposable import Disposable from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In, Out +from dimos.msgs.imitation_msgs.EpisodeStatus import ( + EpisodeEvent, + EpisodeStatus, + RecordingState, +) +from dimos.spec.utils import Spec from dimos.teleop.webxr.controller_types import BUTTON_ALIASES, Buttons from dimos.utils.logging_config import setup_logger @@ -42,25 +47,11 @@ # A button/keyboard press requests one of these; `toggle` resolves to # `start`/`save` based on the current state, so it never reaches the output. EpisodeCommand: TypeAlias = Literal["start", "save", "discard", "toggle"] -# What gets published as `EpisodeStatus.last_event` (`init` on boot). -EpisodeEvent: TypeAlias = Literal["start", "save", "discard", "init"] -RecordingState: TypeAlias = Literal["idle", "recording"] -class EpisodeStatus(BaseModel): - ts: float - state: RecordingState - episodes_saved: int - episodes_discarded: int - last_event: EpisodeEvent = "init" - task_label: str | None = None - - -class KeyPress(BaseModel): - """Single keypress event from a keyboard input source.""" - - key: str - ts: float +class EpisodeControlSpec(Spec, Protocol): + def get_status(self) -> EpisodeStatus: ... + def command(self, event: EpisodeCommand) -> EpisodeStatus: ... def _default_button_map() -> dict[EpisodeCommand, str]: @@ -69,8 +60,14 @@ def _default_button_map() -> dict[EpisodeCommand, str]: class EpisodeMonitorModuleConfig(ModuleConfig): button_map: dict[EpisodeCommand, str] = Field(default_factory=_default_button_map) - keyboard_map: dict[EpisodeCommand, str] = Field(default_factory=dict) - default_task_label: str | None = None + task: str + + @field_validator("task") + @classmethod + def _validate_task(cls, value: str) -> str: + if not value.strip(): + raise ValueError("task must be a non-empty description") + return value.strip() @field_validator("button_map") @classmethod @@ -94,10 +91,7 @@ def _validate_button_map(cls, value: dict[EpisodeCommand, str]) -> dict[EpisodeC class EpisodeMonitorModule(Module): config: EpisodeMonitorModuleConfig - teleop_buttons: In[Buttons] - # TODO: no KeyPress producer exists yet — add a pygame keyboard module that - # publishes KeyPress so this port is actually fed (today only buttons drive it). - keyboard: In[KeyPress] + button_pressed: In[Buttons] status: Out[EpisodeStatus] def __init__(self, **kwargs: Any) -> None: @@ -105,7 +99,7 @@ def __init__(self, **kwargs: Any) -> None: self._state: RecordingState = "idle" self._saved: int = 0 self._discarded: int = 0 - self._prev_bits: dict[str, bool] = {} # rising-edge detection for buttons + self._last_event: EpisodeEvent = "init" self._lock = threading.Lock() self._transition_lock = threading.Lock() self._stopping = False @@ -116,8 +110,7 @@ def start(self) -> None: super().start() # Registered so the base Module.stop() disposes them on shutdown. self._input_subscriptions = [ - self.register_disposable(Disposable(self.teleop_buttons.subscribe(self._on_buttons))), - self.register_disposable(Disposable(self.keyboard.subscribe(self._on_keyboard))), + self.register_disposable(Disposable(self.button_pressed.subscribe(self._on_buttons))), ] # Emit an initial idle status so subscribers (and recorders) have a # known starting point in the timeline. @@ -125,82 +118,58 @@ def start(self) -> None: status = self._snapshot("init", time.time()) self._emit(status) - @rpc - def reset_counters(self) -> EpisodeStatus: - with self._transition_lock: - with self._lock: - if self._stopping: - raise RuntimeError("cannot reset episode counters during shutdown") - self._state = "idle" - self._saved = 0 - self._discarded = 0 - self._prev_bits = {} - status = self._snapshot("init", time.time()) - return self._emit(status) - @rpc def stop(self) -> None: with self._transition_lock: with self._lock: - if self._stopping: - status = None - else: - self._stopping = True - if self._state == "recording": - self._discarded += 1 - self._state = "idle" - status = self._snapshot("discard", time.time()) - else: - status = None + self._stopping = True for subscription in self._input_subscriptions: subscription.dispose() self._input_subscriptions.clear() - if status is not None: - self._emit(status) + # Do not synthesize a save or discard. If the process is interrupted + # mid-take, DataPrep must see the unmatched start as incomplete. super().stop() # ── port handlers ──────────────────────────────────────────────────────── def _on_buttons(self, msg: Buttons) -> None: - """Rising-edge detect against `config.button_map`; advance state machine.""" + """Advance the state machine for configured button-press edges.""" ts = time.time() - # Edge-detect under the lock (it shares `_prev_bits` with reset_counters), - # then fire transitions outside it — `_transition` takes the same lock. fired: list[EpisodeCommand] = [] with self._lock: if self._stopping: return for event_name, alias_or_attr in self.config.button_map.items(): attr = BUTTON_ALIASES.get(alias_or_attr, alias_or_attr) - try: - pressed = bool(getattr(msg, attr)) - except AttributeError: - continue - prev = self._prev_bits.get(attr, False) - self._prev_bits[attr] = pressed - if pressed and not prev: # rising edge + if bool(getattr(msg, attr)): fired.append(event_name) for event_name in fired: self._transition(event_name, ts) - def _on_keyboard(self, msg: KeyPress) -> None: - """Match `msg.key` against `config.keyboard_map`; advance state machine.""" - for event_name, key in self.config.keyboard_map.items(): - if msg.key == key: - self._transition(event_name, msg.ts) - break + @rpc + def command(self, event: EpisodeCommand) -> EpisodeStatus: + """Apply an episode command from an attached operator interface.""" + return self._transition(event, time.time()) - def _transition(self, event: EpisodeCommand, ts: float) -> None: + @rpc + def get_status(self) -> EpisodeStatus: + """Return the latest episode state without publishing a new event.""" + with self._lock: + return self._snapshot(self._last_event, time.time()) + + def _transition(self, event: EpisodeCommand, ts: float) -> EpisodeStatus: """State-machine transition. Publishes EpisodeStatus on every change. ``toggle`` resolves to ``start`` when idle and ``save`` when recording, so one button can begin and end a take. The resolved event is what gets published (DataPrep only ever sees start/save/discard). """ + if event not in ("start", "save", "discard", "toggle"): + raise ValueError(f"unknown episode command: {event!r}") with self._transition_lock: with self._lock: if self._stopping: - return + return self._snapshot(self._last_event, ts) if event == "toggle": event = "save" if self._state == "recording" else "start" if event == "start": @@ -216,9 +185,10 @@ def _transition(self, event: EpisodeCommand, ts: float) -> None: if self._state == "recording": self._discarded += 1 self._state = "idle" + self._last_event = event # Snapshot under the mutation's lock so the event matches the state. status = self._snapshot(event, ts) - self._emit(status) + return self._emit(status) def _snapshot(self, last_event: EpisodeEvent, ts: float) -> EpisodeStatus: """Build a status from current state. Caller must hold `self._lock`.""" @@ -228,7 +198,7 @@ def _snapshot(self, last_event: EpisodeEvent, ts: float) -> EpisodeStatus: episodes_saved=self._saved, episodes_discarded=self._discarded, last_event=last_event, - task_label=self.config.default_task_label, + task_label=self.config.task, ) def _emit(self, status: EpisodeStatus) -> EpisodeStatus: diff --git a/dimos/imitation/collection/profile.py b/dimos/imitation/collection/profile.py new file mode 100644 index 0000000000..e49104e86b --- /dev/null +++ b/dimos/imitation/collection/profile.py @@ -0,0 +1,110 @@ +# 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. + +"""Typed recording inputs and dataset features declared together.""" + +from __future__ import annotations + +from typing import Any + +from pydantic import Field, field_validator, model_validator + +from dimos.imitation.collection.recording import RecordingSchema +from dimos.imitation.dataprep.core import ( + DataPrepConfig, + FeatureSpec, + OutputConfig, + QualityConfig, + SyncConfig, + validate_source_kinds, +) +from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.protocol.service.spec import BaseConfig + + +class CollectionFeature(FeatureSpec): + """A dataset projection and the raw message type needed to record its stream. + + ``message_type`` describes the input, not the exported feature. ``field``, + ``dtype``, and ``shape`` select its dataset representation: for example, + Image data as video or JointState positions converted to float32. + """ + + message_type: type[Any] + + +class CollectionProfile(BaseConfig): + """One recording contract, independent of camera hardware and policy backends.""" + + name: str = Field(min_length=1) + robot_type: str = Field(min_length=1) + observations: dict[str, CollectionFeature] = Field(min_length=1) + actions: dict[str, CollectionFeature] = Field(min_length=1) + sync: SyncConfig + quality: QualityConfig = QualityConfig() + + @field_validator("observations", "actions") + @classmethod + def _feature_names(cls, value: dict[str, CollectionFeature]) -> dict[str, CollectionFeature]: + if any(not key.strip() for key in value): + raise ValueError("feature names must not be blank") + return value + + @model_validator(mode="after") + def validate_features(self) -> CollectionProfile: + if set(self.observations) & set(self.actions): + raise ValueError("observation and action feature names must be distinct") + if self.sync.anchor not in self.observations: + raise ValueError("sync anchor must name an observation feature") + self.input_types() + return self + + def input_types(self) -> dict[str, type[Any]]: + """Deduplicate raw inputs without merging their dataset projections.""" + validate_source_kinds((*self.observations.values(), *self.actions.values())) + inputs: dict[str, type[Any]] = {} + for feature in (*self.observations.values(), *self.actions.values()): + if feature.source_kind == "joint_position_updates" and not issubclass( + feature.message_type, JointState + ): + raise ValueError("joint_position_updates requires JointState messages") + previous = inputs.setdefault(feature.stream, feature.message_type) + if previous is not feature.message_type: + raise ValueError(f"stream {feature.stream!r} has conflicting message types") + return inputs + + def to_schema(self) -> RecordingSchema: + """Snapshot dataset interpretation without runtime message classes.""" + return RecordingSchema( + name=self.name, + robot_type=self.robot_type, + observation={ + key: FeatureSpec(**feature.model_dump(exclude={"message_type"})) + for key, feature in self.observations.items() + }, + action={ + key: FeatureSpec(**feature.model_dump(exclude={"message_type"})) + for key, feature in self.actions.items() + }, + sync=self.sync.model_copy(deep=True), + quality=self.quality.model_copy(deep=True), + ) + + def dataprep_config(self, *, source: str = "", output: OutputConfig) -> DataPrepConfig: + """Build a low-level preparation request for an explicit raw artifact.""" + return DataPrepConfig( + **self.to_schema().model_dump(exclude={"name", "robot_type", "payload"}), + source=source, + output=output, + ) diff --git a/dimos/imitation/collection/recorder.py b/dimos/imitation/collection/recorder.py index 6379464db4..fb643cd2a4 100644 --- a/dimos/imitation/collection/recorder.py +++ b/dimos/imitation/collection/recorder.py @@ -12,41 +12,203 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""CollectionRecorder — captures teleop collection streams to a memory DB. +"""Blueprint factory for recorders with profile-defined typed inputs.""" -A `Recorder` (memory) subscribes each declared `In` port and appends every -message to a SQLite store, flushing durably on stop(). Only *connected* -streams are recorded, so the same recorder works for any arm whose -coordinator publishes `coordinator_joint_state`. +from __future__ import annotations -The recorded stream names match what DataPrep reads: `color_image` -and `coordinator_joint_state` (observation), `status` (episode segmentation). -""" +from abc import ABCMeta +from collections.abc import Callable +import copyreg +import hashlib +import keyword +from pathlib import Path +import sys +from threading import RLock +from typing import Any, ClassVar, Literal, cast, get_type_hints -from __future__ import annotations +from pydantic import Field, field_validator, model_validator +from dimos.core.coordination.blueprints import Blueprint +from dimos.core.core import rpc from dimos.core.stream import In -from dimos.imitation.collection.episode_monitor import EpisodeStatus -from dimos.memory.module import Recorder, RecorderConfig -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.experimental.memory.rust_recorder import ( + RustMcapStoreConfig, + RustRecorder, + RustRecorderConfig, + RustRecordingStoreConfig, + RustSqliteStoreConfig, + RustStreamSpec, +) +from dimos.imitation.collection.profile import CollectionProfile +from dimos.imitation.collection.recording import RecordingSchema +from dimos.memory.module import OnExisting +from dimos.msgs.imitation_msgs.EpisodeStatus import EpisodeStatus + +PortTypes = tuple[tuple[str, type[Any]], ...] -class CollectionRecorderConfig(RecorderConfig): +class CollectionRecorderConfig(RustRecorderConfig): + recording: Path + format: Literal["mcap", "sqlite"] = "mcap" record_tf: bool = False + on_existing: Literal[OnExisting.ERROR] = Field( + default=OnExisting.ERROR, + description="Collection requires a new recording directory; existing files are errors.", + ) + backup_keep_last: Literal[0] = Field( + default=0, description="Collection does not rotate backups." + ) + @field_validator("recording", mode="before") + @classmethod + def _resolve_recording(cls, value: str | Path) -> Path: + return Path(value).expanduser().resolve() -class CollectionRecorder(Recorder): - """Records the streams DataPrep consumes from a teleop session.""" + @model_validator(mode="after") + def _collection_store(self) -> CollectionRecorderConfig: + store: RustRecordingStoreConfig + if self.format == "sqlite": + store = RustSqliteStoreConfig(kind="sqlite", path=str(self.recording / "recording.db")) + else: + store = RustMcapStoreConfig(kind="mcap", path=str(self.recording / "recording.mcap")) + if "store" in self.model_fields_set and self.store != store: + raise ValueError("Collection store is derived from recording and format") + self.store = store + return self + +class CollectionRecorder(RustRecorder): config: CollectionRecorderConfig - color_image: In[Image] # observation (camera) - coordinator_joint_state: In[JointState] # observation + action (measured/next state) - status: In[EpisodeStatus] # episode start/save/discard segmentation + def __init__(self, *, recording_schema: RecordingSchema, **kwargs: Any) -> None: + super().__init__(**kwargs) + self._recording_schema = recording_schema + self._prepared = False + + @rpc + def build(self) -> None: + if self.config.g.replay: + super().build() + return + specs = self._stream_specs() + super().build() + self._prepare_store(specs) + + def _prepare_store(self, specs: list[RustStreamSpec]) -> None: + if self._prepared: + return + schema = self._recording_schema.model_copy(deep=True) + schema.payload = "recording.db" if self.config.format == "sqlite" else "recording.mcap" + streams = {spec.port: spec.name for spec in specs} + for feature in (*schema.observation.values(), *schema.action.values()): + feature.stream = streams[feature.stream] + schema.episodes.status_stream = streams["status"] + encoded = schema.model_dump_json(indent=2) + self.config.recording.mkdir(parents=True, exist_ok=False) + with (self.config.recording / "schema.json").open("x") as file: + file.write(encoded + "\n") + super()._prepare_store(specs) + self._prepared = True + + recording_inputs: ClassVar[PortTypes] = () + + def _stream_specs(self) -> list[RustStreamSpec]: + specs = super()._stream_specs() + connected = {spec.port for spec in specs} + missing = sorted(name for name, _ in self.recording_inputs if name not in connected) + if missing: + raise ValueError(f"Missing required collection inputs: {missing}") + return specs + + +class _RecorderMeta(ABCMeta): + recording_inputs: PortTypes + + +# Class identity must survive standard pickle, fresh workers, and source reload. +# This follows the dynamic relay bridge's metaclass/copyreg protocol. +_class_lock = RLock() +_reserved = ( + frozenset(dir(CollectionRecorder)) + | frozenset(get_type_hints(CollectionRecorder)) + | { + "ref", + "rpc", + "encoded", + "status", + } +) + + +def _recorder_class(ports: PortTypes) -> type[CollectionRecorder]: + with _class_lock: + identity = "\n".join( + f"{name}:{kind.__module__}.{kind.__qualname__}" for name, kind in ports + ) + name = "CollectionRecorder_" + hashlib.sha256(identity.encode()).hexdigest()[:16] + existing = globals().get(name) + if existing is not None: + if existing.recording_inputs != ports: + raise ValueError(f"Recorder class name collision: {name}") + result = cast("type[CollectionRecorder]", existing) + else: + # Runtime types are intentional: each profile supplies these annotations. + input_type: Any = In + result = cast( + "type[CollectionRecorder]", + _RecorderMeta( + name, + (CollectionRecorder,), + { + "__module__": __name__, + "__annotations__": {port: input_type[kind] for port, kind in ports}, + "recording_inputs": ports, + }, + ), + ) + globals()[name] = result + return result + + +def _reduce_recorder(cls: _RecorderMeta) -> tuple[Callable[..., Any], tuple[PortTypes]]: + return _recorder_class, (cls.recording_inputs,) + + +copyreg.pickle(_RecorderMeta, _reduce_recorder) + - async def _resolve_pose(self, name: str, msg: object, ts: float) -> Pose | None: - if name in self.config.poseless_streams: - return None - return await super()._resolve_pose(name, msg, ts) +def collection_recorder( + *, + profile: CollectionProfile, + recording: Path | None = None, + format: Literal["mcap", "sqlite"] = "mcap", + instance_name: str = "recorder", +) -> Blueprint: + """Declare profile inputs and record a portable SQLite or MCAP session directory.""" + inputs = profile.input_types() + for name, kind in inputs.items(): + if ( + not name.isidentifier() + or keyword.iskeyword(name) + or name.startswith("_") + or name in _reserved + ): + raise ValueError(f"Invalid or reserved collection input name: {name!r}") + if ( + kind.__module__ == "__main__" + or kind.__qualname__ != kind.__name__ + or getattr(sys.modules.get(kind.__module__), kind.__name__, None) is not kind + ): + raise ValueError(f"Message type {kind!r} must be importable at module level") + if not callable(getattr(kind, "lcm_encode", None)) or not callable( + getattr(kind, "lcm_decode", None) + ): + raise TypeError(f"Message type {kind!r} does not support native recording") + inputs["status"] = EpisodeStatus + recorder = _recorder_class(tuple(sorted(inputs.items()))) + return recorder.blueprint( + recording_schema=profile.to_schema(), + format=format, + instance_name=instance_name, + **({"recording": recording} if recording is not None else {}), + ) diff --git a/dimos/imitation/collection/recording.py b/dimos/imitation/collection/recording.py new file mode 100644 index 0000000000..7c35af9cbf --- /dev/null +++ b/dimos/imitation/collection/recording.py @@ -0,0 +1,51 @@ +# 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. + +"""Portable collection directories: one raw artifact and its dataset schema.""" + +from pathlib import Path +from typing import Literal + +from dimos.imitation.dataprep.core import DataPrepConfig, DatasetSchema, OutputConfig + + +class RecordingSchema(DatasetSchema): + """JSON snapshot of the collection profile; no Python classes or absolute paths.""" + + name: str + robot_type: str + payload: Literal["recording.mcap", "recording.db"] = "recording.mcap" + + @classmethod + def read(cls, directory: Path) -> "RecordingSchema": + schema = cls.model_validate_json((directory / "schema.json").read_text()) + if not (directory / schema.payload).is_file(): + raise FileNotFoundError(f"Recording payload is missing: {directory / schema.payload}") + return schema + + def dataprep_config(self, directory: Path, output: OutputConfig) -> DataPrepConfig: + output = output.model_copy( + update={ + "metadata": { + "repo_id": f"local/{self.name}", + "robot_type": self.robot_type, + **output.metadata, + } + } + ) + return DataPrepConfig( + **self.model_dump(exclude={"name", "robot_type", "payload"}), + source=str(directory / self.payload), + output=output, + ) diff --git a/dimos/imitation/collection/test_blueprint.py b/dimos/imitation/collection/test_blueprint.py index 43e352dd10..f901c9a566 100644 --- a/dimos/imitation/collection/test_blueprint.py +++ b/dimos/imitation/collection/test_blueprint.py @@ -33,15 +33,22 @@ "blueprint", [learning_collect_webxr_xarm7, learning_collect_webxr_piper], ) -def test_collection_streams_are_poseless(blueprint: Blueprint) -> None: - recorder = next(atom for atom in blueprint.blueprints if atom.module is CollectionRecorder) - - assert recorder.kwargs["poseless_streams"] == [ - "color_image", - "coordinator_joint_state", - "status", - ] - assert recorder.kwargs["record_tf"] is False +def test_collection_uses_sqlite_without_tf(blueprint: Blueprint) -> None: + recorder = next( + atom for atom in blueprint.blueprints if issubclass(atom.module, CollectionRecorder) + ) + + assert recorder.kwargs["format"] == "sqlite" + schema = recorder.kwargs["recording_schema"] + assert schema.action["action"] == schema.observation["state"] + assert schema.observation["camera"].shape == (480, 848, 3) + assert schema.sync.rate_hz == 15 + coordinator = next( + atom for atom in blueprint.active_blueprints if atom.name == "ControlCoordinator" + ) + joints = [joint for hardware in coordinator.kwargs["hardware"] for joint in hardware.joints] + assert schema.observation["state"].names == joints + assert schema.observation["state"].shape == (len(joints),) @pytest.mark.parametrize( @@ -49,7 +56,7 @@ def test_collection_streams_are_poseless(blueprint: Blueprint) -> None: [learning_collect_webxr_xarm7, learning_collect_webxr_piper], ) def test_collection_recorder_stops_after_producers(blueprint: Blueprint) -> None: - assert blueprint.active_blueprints[0].module is CollectionRecorder + assert issubclass(blueprint.active_blueprints[0].module, CollectionRecorder) @pytest.mark.parametrize( @@ -82,12 +89,14 @@ def _joint_streams(blueprint: Blueprint) -> dict[tuple[str, str], str]: } -@pytest.mark.parametrize("blueprint", [learning_collect_webxr_xarm7, learning_collect_webxr_piper]) +@pytest.mark.parametrize( + "blueprint", + [learning_collect_webxr_xarm7, learning_collect_webxr_piper], +) def test_recorder_reads_aggregate_joint_state(blueprint: Blueprint) -> None: streams = _joint_streams(blueprint) # Plain name pairing on both ends, no remap in between. The coordinator # atom carries its explicit instance_name (the RPC lookup contract). - assert streams[("collectionrecorder", AGGREGATE)] == AGGREGATE + assert streams[("recorder", AGGREGATE)] == AGGREGATE assert streams[("ControlCoordinator", AGGREGATE)] == AGGREGATE - assert not [port for _instance, port in streams if port.endswith("_joints")] diff --git a/dimos/imitation/collection/test_episode_monitor.py b/dimos/imitation/collection/test_episode_monitor.py index fb0549c9d2..dc171d3735 100644 --- a/dimos/imitation/collection/test_episode_monitor.py +++ b/dimos/imitation/collection/test_episode_monitor.py @@ -30,11 +30,8 @@ import pytest import pytest_mock -from dimos.imitation.collection.episode_monitor import ( - EpisodeMonitorModule, - EpisodeStatus, - KeyPress, -) +from dimos.imitation.collection.episode_monitor import EpisodeMonitorModule +from dimos.msgs.imitation_msgs.EpisodeStatus import EpisodeStatus from dimos.protocol.rpc.pubsubrpc import LCMRPC from dimos.teleop.webxr.controller_types import BUTTON_ALIASES, Buttons @@ -59,6 +56,7 @@ def make_monitor( built: list[EpisodeMonitorModule] = [] def _make(**config: object) -> EpisodeMonitorModule: + config.setdefault("task", "pick up the block") m = EpisodeMonitorModule(**config) m.status = mocker.MagicMock() # type: ignore[assignment] built.append(m) @@ -75,12 +73,10 @@ def _events(monitor: EpisodeMonitorModule) -> list[EpisodeStatus]: def _press(monitor: EpisodeMonitorModule, alias: str) -> None: - """Rising edge: release-then-press the given WebXR button alias.""" + """Deliver one debounced WebXR button-press edge.""" attr = BUTTON_ALIASES[alias] - released = Buttons() pressed = Buttons() pressed.set_attribute(attr, True) - monitor._on_buttons(released) monitor._on_buttons(pressed) @@ -94,6 +90,52 @@ def test_toggle_starts_then_saves(make_monitor: Callable[..., EpisodeMonitorModu assert events[-1].state == "idle" assert events[-1].episodes_saved == 1 assert events[-1].episodes_discarded == 0 + assert events[-1].task_label == "pick up the block" + + +def test_rpc_commands_use_the_same_state_machine( + make_monitor: Callable[..., EpisodeMonitorModule], +) -> None: + m = make_monitor() + + recording = m.command("toggle") + saved = m.command("toggle") + + assert [event.last_event for event in _events(m)] == ["start", "save"] + assert recording.state == "recording" + assert saved.state == "idle" + assert saved.episodes_saved == 1 + + +def test_get_status_does_not_publish( + make_monitor: Callable[..., EpisodeMonitorModule], +) -> None: + m = make_monitor() + m.command("start") + event_count = len(_events(m)) + + status = m.get_status() + + assert status.state == "recording" + assert status.last_event == "start" + assert len(_events(m)) == event_count + + +def test_invalid_rpc_command_is_rejected_without_changing_state( + make_monitor: Callable[..., EpisodeMonitorModule], +) -> None: + m = make_monitor() + + with pytest.raises(ValueError, match="unknown episode command"): + m.command("pause") # type: ignore[arg-type] + + assert m.get_status().state == "idle" + assert _events(m) == [] + + +def test_task_is_required(make_monitor: Callable[..., EpisodeMonitorModule]) -> None: + with pytest.raises(ValidationError, match="task"): + EpisodeMonitorModule() def test_discard_does_not_count_as_saved( @@ -112,11 +154,9 @@ def test_discard_does_not_count_as_saved( def test_start_while_recording_autocommits_previous( make_monitor: Callable[..., EpisodeMonitorModule], ) -> None: - # toggle (start), then an explicit start via keyboard while still recording: - # the in-progress episode auto-commits (matches the offline extractor). - m = make_monitor(keyboard_map={"start": "r"}) - _press(m, "B") # recording - m._on_keyboard(KeyPress(key="r", ts=2.0)) # start again → auto-commit prior + m = make_monitor(button_map={"start": "A"}) + _press(m, "A") + _press(m, "A") last = _events(m)[-1] assert last.last_event == "start" @@ -124,17 +164,6 @@ def test_start_while_recording_autocommits_previous( assert last.episodes_saved == 1 # the auto-committed one -def test_no_event_without_rising_edge( - make_monitor: Callable[..., EpisodeMonitorModule], -) -> None: - m = make_monitor() - pressed = Buttons() - pressed.right_secondary = True # B held - m._on_buttons(pressed) - m._on_buttons(pressed) # still held — no new edge - assert [e.last_event for e in _events(m)] == ["start"] - - def test_published_status_is_internally_consistent( make_monitor: Callable[..., EpisodeMonitorModule], ) -> None: @@ -158,27 +187,18 @@ def test_published_status_is_internally_consistent( assert events[-1].episodes_discarded == 1 -def test_reset_counters(make_monitor: Callable[..., EpisodeMonitorModule]) -> None: - m = make_monitor() - _press(m, "B") - _press(m, "B") - status = m.reset_counters() - assert status.episodes_saved == 0 - assert status.episodes_discarded == 0 - assert status.state == "idle" - assert status.last_event == "init" - - -def test_shutdown_discards_recording(make_monitor: Callable[..., EpisodeMonitorModule]) -> None: +def test_shutdown_leaves_recording_incomplete( + make_monitor: Callable[..., EpisodeMonitorModule], +) -> None: m = make_monitor() _press(m, "B") m.stop() last = _events(m)[-1] - assert last.last_event == "discard" - assert last.state == "idle" - assert last.episodes_discarded == 1 + assert last.last_event == "start" + assert last.state == "recording" + assert last.episodes_discarded == 0 def test_invalid_button_mapping_fails_at_startup( @@ -204,8 +224,6 @@ def test_buttons_are_ignored_after_shutdown_begins( _press(m, "B") assert _events(m) == [] - with pytest.raises(RuntimeError, match="during shutdown"): - m.reset_counters() def test_stop_waits_for_in_flight_transition_and_blocks_later_transitions( diff --git a/dimos/imitation/collection/test_profile.py b/dimos/imitation/collection/test_profile.py new file mode 100644 index 0000000000..91b612daca --- /dev/null +++ b/dimos/imitation/collection/test_profile.py @@ -0,0 +1,114 @@ +# 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. + + +import numpy as np +import pytest + +from dimos.imitation.collection.profile import CollectionFeature, CollectionProfile +from dimos.imitation.dataprep.core import DataPrepConfig, OutputConfig, SyncConfig, resolve_field +from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.sensor_msgs.JointState import JointState + + +@pytest.fixture +def profile(): + return CollectionProfile( + name="custom", + robot_type="test", + observations={ + "camera": CollectionFeature( + stream="camera", + message_type=Image, + field="data", + dtype="video", + shape=(8, 8, 3), + names=["height", "width", "channels"], + ), + "positions": CollectionFeature( + stream="measured", + message_type=JointState, + field="position", + dtype="float32", + shape=(2,), + names=["b", "a"], + ), + "velocities": CollectionFeature( + stream="measured", + message_type=JointState, + field="velocity", + dtype="float32", + shape=(1,), + names=["a"], + ), + }, + actions={ + "target": CollectionFeature( + stream="commanded", + source_kind="joint_position_updates", + message_type=JointState, + field="position", + dtype="float32", + shape=(1,), + names=["b"], + ), + }, + sync=SyncConfig(anchor="camera", rate_hz=30, tolerance_ms=20), + ) + + +def test_profile_lowers_to_existing_json_protocol_and_preserves_projections(profile, tmp_path): + config = profile.dataprep_config(output=OutputConfig(path=tmp_path)) + restored = DataPrepConfig.model_validate_json(config.model_dump_json()) + assert profile.input_types() == { + "camera": Image, + "measured": JointState, + "commanded": JointState, + } + assert restored.observation["positions"].names == ["b", "a"] + assert restored.action["target"].names == ["b"] + assert restored.action["target"].source_kind == "joint_position_updates" + assert restored.observation["positions"].source_kind == "snapshot" + message = JointState(name=["a", "b"], position=[1.0, 2.0], velocity=[3.0, 4.0]) + np.testing.assert_array_equal(resolve_field(message, restored.observation["positions"]), [2, 1]) + np.testing.assert_array_equal(resolve_field(message, restored.observation["velocities"]), [3]) + + +def test_profile_rejects_conflicting_raw_types(profile): + values = profile.model_dump() + values["actions"]["target"].update(stream="camera", source_kind="snapshot") + with pytest.raises(ValueError, match="conflicting message types"): + CollectionProfile(**values) + + +def test_profile_rejects_conflicting_source_kinds(profile): + values = profile.model_dump() + values["actions"]["target"].update(stream="measured", source_kind="joint_position_updates") + with pytest.raises(ValueError, match="conflicting source kinds"): + CollectionProfile(**values) + + +def test_profile_rejects_updates_from_non_joint_messages(profile): + values = profile.model_dump() + values["actions"]["target"].update(message_type=Image, source_kind="joint_position_updates") + with pytest.raises(ValueError, match="requires JointState"): + CollectionProfile(**values) + + +@pytest.mark.parametrize("anchor", ["absent", "target"]) +def test_sync_anchor_must_be_an_observation(profile, anchor): + values = profile.model_dump() + values["sync"]["anchor"] = anchor + with pytest.raises(ValueError, match="sync anchor"): + CollectionProfile(**values) diff --git a/dimos/imitation/collection/test_recorder.py b/dimos/imitation/collection/test_recorder.py index b854c23f94..624b564d26 100644 --- a/dimos/imitation/collection/test_recorder.py +++ b/dimos/imitation/collection/test_recorder.py @@ -12,17 +12,425 @@ # See the License for the specific language governing permissions and # limitations under the License. -import pytest_mock +import importlib +from pathlib import Path +import pickle +import sys -from dimos.imitation.collection.recorder import CollectionRecorder, CollectionRecorderConfig +import pytest +from dimos.core.coordination.blueprint_config.parser import BlueprintConfigParser +from dimos.core.coordination.blueprints import autoconnect +from dimos.core.coordination.worker_manager_python import WorkerManagerPython +from dimos.core.global_config import GlobalConfig, global_config +from dimos.core.module import Module +from dimos.core.native_module import NativeModule +from dimos.core.stream import Out, RemoteIn +from dimos.imitation.collection import recorder as recorder_module +from dimos.imitation.collection.profile import CollectionFeature, CollectionProfile +from dimos.imitation.collection.recorder import CollectionRecorderConfig, collection_recorder +from dimos.imitation.collection.recording import RecordingSchema +from dimos.imitation.dataprep.core import OutputConfig, SyncConfig +from dimos.msgs.imitation_msgs.EpisodeStatus import EpisodeStatus +from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.robot.get_all_blueprints import get_by_name -async def test_poseless_collection_stream_skips_pose_lookup( - mocker: pytest_mock.MockerFixture, -) -> None: - recorder = mocker.MagicMock(spec=CollectionRecorder) - recorder.config = CollectionRecorderConfig(poseless_streams=["commands"]) - pose = await CollectionRecorder._resolve_pose(recorder, "commands", object(), 1.0) +class _CameraPublisher(Module): + color_image: Out[Image] - assert pose is None + +def _profile(camera_count): + return CollectionProfile( + name="test", + robot_type="test", + observations={ + **{ + f"images.{i}": CollectionFeature( + stream=f"camera_{i}", + message_type=Image, + field="data", + dtype="video", + shape=(8, 8, 3), + names=["height", "width", "channels"], + ) + for i in range(camera_count) + }, + "state": CollectionFeature( + stream="measured", + message_type=JointState, + field="position", + dtype="float32", + shape=(1,), + names=["joint"], + ), + }, + actions={ + "action": CollectionFeature( + stream="measured", + message_type=JointState, + field="position", + dtype="float32", + shape=(1,), + names=["joint"], + ) + }, + sync=SyncConfig(anchor="images.0", rate_hz=30, tolerance_ms=20), + ) + + +@pytest.fixture +def recorder(tmp_path): + profile = _profile(2) + profile.actions["action"].stream = "commanded" + atom = collection_recorder( + profile=profile, + recording=tmp_path / "collection.mcap", + ).active_blueprints[0] + instance = atom.module(**atom.kwargs) + yield instance + instance.stop() + + +@pytest.mark.parametrize("count", [1, 2, 4]) +def test_factory_exposes_all_typed_inputs_before_autoconnect(count, tmp_path): + blueprint = collection_recorder(profile=_profile(count), recording=tmp_path / "test.mcap") + stack = autoconnect( + blueprint, + *[_CameraPublisher.blueprint(instance_name=f"publisher_{i}") for i in range(count)], + ).remappings([(f"publisher_{i}", "color_image", f"camera_{i}") for i in range(count)]) + atom = stack.active_blueprints[0] + assert {s.name: s.type for s in atom.streams if s.name != "tf"} == { + **{f"camera_{i}": Image for i in range(count)}, + "measured": JointState, + "status": EpisodeStatus, + } + assert len([s for s in atom.streams if s.name == "measured"]) == 1 + # Autoconnect groups remapped ports by (name, message type). + for i in range(count): + peers = [ + (peer.name, port.direction) + for peer in stack.active_blueprints + for port in peer.streams + if stack.remapping_map.get((peer.name, port.name), port.name) == f"camera_{i}" + and port.type is Image + ] + assert peers == [(atom.name, "in"), (f"publisher_{i}", "out")] + + +def test_two_shapes_and_recording_paths_have_independent_configuration(tmp_path): + first = collection_recorder( + profile=_profile(1), recording=tmp_path / "one.mcap", instance_name="one" + ) + second = collection_recorder( + profile=_profile(4), recording=tmp_path / "four.mcap", instance_name="four" + ) + same_shape = collection_recorder(profile=_profile(1), recording=tmp_path / "other.mcap") + one, four = autoconnect(first, second).active_blueprints + assert one.module is same_shape.active_blueprints[0].module + assert one.module is not four.module + assert one.kwargs["recording"] == tmp_path / "one.mcap" + assert four.kwargs["recording"] == tmp_path / "four.mcap" + + +def test_blueprint_pickle_and_reload_preserve_generated_class(tmp_path): + blueprint = collection_recorder(profile=_profile(2), recording=tmp_path / "recording.mcap") + module = blueprint.active_blueprints[0].module + payload = pickle.dumps(blueprint) + assert pickle.loads(payload).active_blueprints[0].module is module + importlib.reload(recorder_module) + restored = pickle.loads(payload) + assert restored.active_blueprints[0].module is module + assert pickle.loads(pickle.dumps(restored)).active_blueprints[0].module is module + rebuilt = module.blueprint(**restored.active_blueprints[0].kwargs) + assert rebuilt.active_blueprints[0].streams == restored.active_blueprints[0].streams + + +def test_recorder_resolves_every_connected_stream(recorder, mocker): + for name in ( + "camera_0", + "camera_1", + "measured", + "commanded", + "status", + ): + getattr(recorder, name).transport = mocker.MagicMock(channel=f"dimos/{name}") + assert {s.name: s.codec for s in recorder._stream_specs()} == { + "camera_0": "jpeg", + "camera_1": "jpeg", + "measured": "lcm", + "commanded": "lcm", + "status": "lcm", + } + + +def test_missing_input_fails_before_native_process_starts(recorder, mocker): + start = mocker.patch("dimos.experimental.memory.rust_recorder.NativeModule.start") + with pytest.raises(ValueError, match="Missing required collection inputs"): + recorder.start() + start.assert_not_called() + + +@pytest.mark.parametrize( + "stream", ["status", "tf", "start", "config", "rpc", "bad-name", "_private"] +) +def test_invalid_ports_fail_at_factory_boundary(stream, tmp_path): + profile = _profile(1) + profile.observations["images.0"].stream = stream + with pytest.raises(ValueError, match="reserved"): + collection_recorder(profile=profile, recording=tmp_path / "invalid.mcap") + + +def test_native_collection_uses_the_recorder_build_directory(recorder): + root = Path(__file__).parents[2] / "experimental" / "memory" / "rust" + assert Path(recorder.config.cwd) == root + assert Path(recorder.config.executable) == root / "result/bin/dimos-memory-recorder" + + +@pytest.fixture +def worker_manager(): + manager = WorkerManagerPython(g=GlobalConfig(n_workers=1)) + manager.start() + yield manager + manager.stop() + + +@pytest.fixture +def deployed_recorders(worker_manager): + proxies = [] + yield proxies + for proxy in reversed(proxies): + proxy.stop() + + +@pytest.mark.skipif_macos_bug +def test_generated_inputs_survive_forkserver_and_fresh_deployment( + worker_manager, + deployed_recorders, + tmp_path, +): + # Workers predate the class; fork inheritance cannot make this pass. + assert worker_manager.workers[0].pid is not None + atom = collection_recorder( + profile=_profile(3), recording=tmp_path / "worker.mcap" + ).active_blueprints[0] + first = worker_manager.deploy( + atom.module, global_config, {**atom.kwargs, "instance_name": "first"} + ) + deployed_recorders.append(first) + importlib.reload(recorder_module) + reloaded = getattr(recorder_module, atom.module.__name__) + fresh = worker_manager.deploy_fresh( + reloaded, global_config, {**atom.kwargs, "instance_name": "fresh"} + ) + deployed_recorders.append(fresh) + for proxy in (first, fresh): + for name, kind in atom.module.recording_inputs: + port = getattr(proxy, name) + assert isinstance(port, RemoteIn) + assert port.type is kind + pids = [worker.pid for worker in worker_manager.workers] + assert len(set(pids)) == 2 + + +def test_unsupported_message_type_is_rejected(tmp_path): + profile = _profile(1) + profile.observations["images.0"].message_type = str + with pytest.raises(TypeError, match="native recording"): + collection_recorder(profile=profile, recording=tmp_path / "invalid.mcap") + + +def test_local_message_type_is_rejected(tmp_path): + class LocalMessage: + pass + + profile = _profile(1) + profile.observations["images.0"].message_type = LocalMessage + with pytest.raises(ValueError, match="importable at module level"): + collection_recorder(profile=profile, recording=tmp_path / "invalid.mcap") + + +def test_run_config_resolves_collection_destination(tmp_path): + blueprint = collection_recorder(profile=_profile(1)) + parser = BlueprintConfigParser(blueprint) + help_text = parser.format_help() + assert "--recorder.recording" in help_text + assert "--recorder.format" in help_text + assert "recording-schema" not in help_text + assert "store.path" in help_text + parsed = parser.parse( + ["--recorder.recording", str(tmp_path / "session"), "--recorder.format", "sqlite"], + environ={}, + ) + atom = blueprint.active_blueprints[0] + recorder = atom.module(**{**atom.kwargs, **parsed.module_kwargs(atom.name)}) + try: + assert recorder.config.store.path == str(tmp_path / "session" / "recording.db") + assert recorder._recording_schema.observation["images.0"].stream == "camera_0" + finally: + recorder.stop() + + +def test_same_ports_keep_independent_dataset_projections(tmp_path): + first_profile = _profile(1) + second_profile = _profile(1) + second_profile.observations["state"].names = ["other_joint"] + first = collection_recorder( + profile=first_profile, recording=tmp_path / "first" + ).active_blueprints[0] + second = collection_recorder( + profile=second_profile, recording=tmp_path / "second" + ).active_blueprints[0] + assert first.module is second.module + assert first.kwargs["recording_schema"].observation["state"].names == ["joint"] + assert second.kwargs["recording_schema"].observation["state"].names == ["other_joint"] + + +@pytest.fixture +def connected_recorder(tmp_path, mocker): + def make(format="mcap", **kwargs): + atom = collection_recorder( + profile=_profile(2), recording=tmp_path / "session", format=format + ).active_blueprints[0] + instance = atom.module(**atom.kwargs, **kwargs) + for port, _ in instance.recording_inputs: + getattr(instance, port).transport = mocker.MagicMock(channel=f"dimos/{port}") + recorders.append(instance) + return instance + + recorders = [] + yield make + for instance in recorders: + instance.stop() + + +@pytest.mark.parametrize( + ("format", "payload"), [("mcap", "recording.mcap"), ("sqlite", "recording.db")] +) +def test_build_saves_portable_schema_before_native_capture( + connected_recorder, format, payload, mocker +): + recorder = connected_recorder( + format, stream_remapping={"camera_0": "wrist", "status": "episodes"} + ) + mocker.patch.object(NativeModule, "build") + start = mocker.patch.object(NativeModule, "start") + recorder.build() + directory = recorder.config.recording + schema = RecordingSchema.model_validate_json((directory / "schema.json").read_text()) + assert schema.payload == payload + assert schema.observation["images.0"].stream == "wrist" + assert schema.episodes.status_stream == "episodes" + assert schema.action["action"].names == ["joint"] + config = schema.dataprep_config(directory, OutputConfig(path=directory.parent / "dataset")) + assert config.source == str(directory / payload) + start.assert_not_called() + recorder.start() + start.assert_called_once_with() + assert recorder.config.to_config_dict()["store"]["path"] == str(directory / payload) + + +def test_existing_directory_is_never_overwritten(connected_recorder, mocker): + recorder = connected_recorder() + recorder.config.recording.mkdir() + marker = recorder.config.recording / "schema.json" + marker.write_text("existing") + mocker.patch.object(NativeModule, "build") + start = mocker.patch.object(NativeModule, "start") + with pytest.raises(FileExistsError): + recorder.build() + assert marker.read_text() == "existing" + start.assert_not_called() + + +def test_missing_connections_fail_before_build_or_directory_creation(recorder, mocker): + native_build = mocker.patch.object(NativeModule, "build") + with pytest.raises(ValueError, match="Missing required collection inputs"): + recorder.build() + native_build.assert_not_called() + assert not recorder.config.recording.exists() + + +def test_schema_write_failure_prevents_capture(connected_recorder, mocker): + recorder = connected_recorder() + mocker.patch.object(NativeModule, "build") + start = mocker.patch.object(NativeModule, "start") + mocker.patch.object(Path, "open", side_effect=PermissionError("not writable")) + with pytest.raises(PermissionError, match="not writable"): + recorder.build() + start.assert_not_called() + assert not recorder._prepared + + +def test_external_package_uses_standard_blueprint_entrypoint(tmp_path, monkeypatch): + package = tmp_path / "vendor_robot" + package.mkdir() + (package / "__init__.py").write_text("") + (package / "collection.py").write_text( + "from dimos.core.coordination.blueprints import autoconnect\n" + "from dimos.imitation.collection.recorder import CollectionRecorderConfig, collection_recorder\n" + "from dimos.imitation.collection.profile import CollectionFeature, CollectionProfile\n" + "from dimos.imitation.dataprep.core import SyncConfig\n" + "from dimos.msgs.sensor_msgs.JointState import JointState\n" + "feature = CollectionFeature(stream='joints', message_type=JointState, field='position', dtype='float32', shape=(1,), names=['joint'])\n" + "profile = CollectionProfile(name='vendor', robot_type='vendor', observations={'state': feature}, actions={'action': feature}, sync=SyncConfig(anchor='state', rate_hz=30, tolerance_ms=20))\n" + "collect = autoconnect(collection_recorder(profile=profile))\n" + ) + metadata = tmp_path / "vendor_robot-1.0.dist-info" + metadata.mkdir() + (metadata / "METADATA").write_text("Metadata-Version: 2.1\nName: vendor-robot\nVersion: 1.0\n") + (metadata / "entry_points.txt").write_text( + "[dimos.blueprints]\ncollect = vendor_robot.collection:collect\n" + ) + monkeypatch.syspath_prepend(str(tmp_path)) + try: + blueprint = get_by_name("vendor-robot.collect") + parsed = BlueprintConfigParser(blueprint).parse( + ["--recording", str(tmp_path / "session")], environ={} + ) + assert parsed.module_kwargs("recorder")["recording"] == tmp_path / "session" + assert blueprint.active_blueprints[0].kwargs["recording_schema"].robot_type == "vendor" + assert {port.name for port in blueprint.active_blueprints[0].streams} >= { + "joints", + "status", + } + finally: + sys.modules.pop("vendor_robot.collection", None) + sys.modules.pop("vendor_robot", None) + + +@pytest.mark.parametrize("format,payload", [("mcap", "recording.mcap"), ("sqlite", "recording.db")]) +def test_collection_config_roundtrip_keeps_derived_store(tmp_path, format, payload): + config = CollectionRecorderConfig(recording=tmp_path / "session", format=format) + restored = CollectionRecorderConfig.model_validate(config.model_dump()) + assert restored.store.path == str(tmp_path / "session" / payload) + assert restored.to_config_dict() == config.to_config_dict() + + +@pytest.mark.parametrize( + "kwargs", + [ + {"store": {"kind": "sqlite", "path": "other.db"}}, + {"store": {"kind": "mcap", "path": "other.mcap"}}, + {"on_existing": "overwrite"}, + {"on_existing": "backup"}, + {"on_existing": "append"}, + {"backup_keep_last": 10}, + ], +) +def test_collection_rejects_conflicting_file_settings_before_creating_directory(tmp_path, kwargs): + directory = tmp_path / "session" + with pytest.raises(ValueError): + CollectionRecorderConfig(recording=directory, **kwargs) + assert not directory.exists() + + +def test_replay_does_not_prepare_collection(connected_recorder, mocker): + recorder = connected_recorder(g=GlobalConfig(replay=True)) + mocker.patch.object(NativeModule, "build") + native_start = mocker.patch.object(NativeModule, "start") + recorder.build() + recorder.start() + native_start.assert_not_called() + assert not recorder.config.recording.exists() diff --git a/dimos/imitation/collection/test_recording.py b/dimos/imitation/collection/test_recording.py new file mode 100644 index 0000000000..dc73754f71 --- /dev/null +++ b/dimos/imitation/collection/test_recording.py @@ -0,0 +1,32 @@ +# 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. + + +import pytest + +from dimos.imitation.collection.recording import RecordingSchema + + +@pytest.mark.parametrize("payload", ["../elsewhere.db", "/tmp/elsewhere.mcap"]) +def test_schema_cannot_redirect_preparation_outside_directory(payload): + with pytest.raises(ValueError, match="payload"): + RecordingSchema(name="test", robot_type="test", payload=payload) + + +def test_missing_payload_is_reported(tmp_path): + (tmp_path / "schema.json").write_text( + RecordingSchema(name="test", robot_type="test").model_dump_json() + ) + with pytest.raises(FileNotFoundError, match="payload is missing"): + RecordingSchema.read(tmp_path) diff --git a/dimos/imitation/dataprep/_lerobot_protocol.py b/dimos/imitation/dataprep/_lerobot_protocol.py new file mode 100644 index 0000000000..a5e3c9ee31 --- /dev/null +++ b/dimos/imitation/dataprep/_lerobot_protocol.py @@ -0,0 +1,50 @@ +# 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. + +"""Private typed wire shared by the host and isolated LeRobot process.""" + +from pathlib import Path +from typing import Annotated, Any, Literal, TypeAlias + +from pydantic import BaseModel, Field, TypeAdapter + +from dimos.imitation.dataprep.core import DataPrepConfig + + +class BuildRequest(BaseModel): + command: Literal["build"] = "build" + config: DataPrepConfig + + +class InspectRequest(BaseModel): + command: Literal["inspect"] = "inspect" + path: Path + + +Request: TypeAlias = Annotated[BuildRequest | InspectRequest, Field(discriminator="command")] +REQUEST_ADAPTER: TypeAdapter[Request] = TypeAdapter(Request) + + +class BuildResult(BaseModel): + command: Literal["build"] = "build" + path: Path + + +class InspectResult(BaseModel): + command: Literal["inspect"] = "inspect" + info: dict[str, Any] + + +Result: TypeAlias = Annotated[BuildResult | InspectResult, Field(discriminator="command")] +RESULT_ADAPTER: TypeAdapter[Result] = TypeAdapter(Result) diff --git a/dimos/imitation/dataprep/build.py b/dimos/imitation/dataprep/build.py index f3e000184c..0179e4ce9c 100644 --- a/dimos/imitation/dataprep/build.py +++ b/dimos/imitation/dataprep/build.py @@ -17,34 +17,79 @@ `run_dataprep` (build) and `inspect_dataset` (read-back) own the I/O and side effects — open/close the store, drive the writer/reader, emit logs, write files; they compose the pure helpers in `core.py` and the per-format -readers/writers. Exposed by the `dimos dataprep` subcommand. +readers/writers. Built-in workflows expose it through `dimos imitation`. """ from __future__ import annotations from collections.abc import Iterator +from itertools import chain import json from pathlib import Path -from typing import Any +from typing import Any, cast + +from mcap.reader import make_reader from dimos.imitation.dataprep.core import ( DataPrepConfig, Episode, EpisodeExtractor, + EpisodeQualityReport, Sample, + Writer, extract_episodes, - get_inspector, get_writer, + inspect_episode_quality, inspect_episodes, iter_episode_samples, ) +from dimos.memory.codecs.base import codec_from_id +from dimos.memory.store.base import Store +from dimos.memory.store.mcap import McapStore, StreamCodec from dimos.memory.store.sqlite import SqliteStore from dimos.utils.logging_config import setup_logger logger = setup_logger() -def _write_dimos_meta(dataset_path: Path, config: DataPrepConfig, episodes: list[Episode]) -> None: +def _open_recording(path: str | Path) -> Store: + """Open a native collection artifact through its read-only Store interface.""" + source = Path(path) + if source.suffix == ".db": + return SqliteStore(path=str(source), must_exist=True) + if source.suffix == ".mcap": + return McapStore(path=str(source), codecs=_recording_codecs(source)) + raise ValueError(f"Unsupported recording {str(source)!r}: expected a .db or .mcap artifact") + + +def _recording_codecs(path: Path) -> dict[str, StreamCodec]: + """Load native codecs from a trusted recording's message-type metadata.""" + with path.open("rb") as file: + summary = make_reader(file).get_summary() + codecs: dict[str, StreamCodec] = {} + if summary is None: + return codecs + for channel in summary.channels.values(): + payload_type = channel.metadata.get("dimos.payload_type") + if payload_type and channel.message_encoding in {"jpeg", "lcm", "lz4+lcm"}: + try: + codecs[channel.topic] = cast( + "StreamCodec", codec_from_id(channel.message_encoding, payload_type) + ) + except (ImportError, AttributeError) as exc: + raise ImportError( + f"Cannot decode MCAP stream {channel.topic!r}: install the package " + f"providing {payload_type!r} in the dataset reader environment" + ) from exc + return codecs + + +def _write_dimos_meta( + dataset_path: Path, + config: DataPrepConfig, + episodes: list[Episode], + quality_reports: list[EpisodeQualityReport], +) -> None: """Sidecar describing how this dataset was built, recording the obs/action schema alongside the dataset.""" meta = { @@ -52,6 +97,8 @@ def _write_dimos_meta(dataset_path: Path, config: DataPrepConfig, episodes: list "observation": {k: v.model_dump() for k, v in config.observation.items()}, "action": {k: v.model_dump() for k, v in config.action.items()}, "sync": config.sync.model_dump(), + "quality": config.quality.model_dump(), + "quality_reports": [report.model_dump() for report in quality_reports], "episodes": [ { "id": e.id, @@ -75,7 +122,7 @@ def _write_dimos_meta(dataset_path: Path, config: DataPrepConfig, episodes: list json.dump(meta, f, indent=2, default=str) -def run_dataprep(config: DataPrepConfig) -> Path: +def run_dataprep(config: DataPrepConfig, *, writer: Writer | None = None) -> Path: """Build a dataset from a recording and return the dataset path. Opens the source store, extracts episodes, streams samples through the @@ -95,7 +142,7 @@ def run_dataprep(config: DataPrepConfig) -> Path: config.episodes.extractor, config.output.path, ) - store = SqliteStore(path=config.source, must_exist=True) + store = _open_recording(config.source) try: logger.info("[dataprep] streams in source: %s", store.list_streams()) all_eps = extract_episodes(store, config.episodes) @@ -124,10 +171,42 @@ def run_dataprep(config: DataPrepConfig) -> Path: sorted(action_keys), config.sync.model_dump(), ) - writer = get_writer(config.output.format) + selected_writer = writer or get_writer(config.output.format) + quality_reports = [ + inspect_episode_quality(store, episode, streams, config.sync, config.quality) + for episode in successful + ] + valid_ids = {report.episode_id for report in quality_reports if report.valid} + valid_episodes = [episode for episode in successful if episode.id in valid_ids] + for report in quality_reports: + if not report.valid: + logger.warning( + "[dataprep] excluding episode %s: %s", + report.episode_id, + "; ".join(report.rejection_reasons), + ) + if not valid_episodes: + reasons = {report.episode_id: report.rejection_reasons for report in quality_reports} + raise RuntimeError(f"All saved episodes failed dataset validation: {reasons}") # fps drives written timestamps + video rate, so tie it to the resample # rate; an explicit metadata.fps still wins. - output = config.output + feature_schema = { + **{key: value.model_dump() for key, value in config.observation.items()}, + **{key: value.model_dump() for key, value in config.action.items()}, + "complementary_info.is_filled": { + "dtype": "bool", + "shape": [1], + "names": ["is_filled"], + }, + } + output = config.output.model_copy( + update={ + "metadata": { + **config.output.metadata, + "feature_schema": feature_schema, + } + } + ) if config.sync.rate_hz > 0 and "fps" not in output.metadata: output = output.model_copy( update={"metadata": {**output.metadata, "fps": config.sync.rate_hz}} @@ -136,18 +215,19 @@ def run_dataprep(config: DataPrepConfig) -> Path: samples_seen = 0 episodes_done = 0 - total = len(successful) + total = len(valid_episodes) produced: list[Episode] = [] # episodes that yielded ≥1 sample def _all_samples() -> Iterator[Sample]: nonlocal samples_seen, episodes_done - for ep in successful: + for ep in valid_episodes: before = samples_seen for sample in iter_episode_samples( store=store, episode=ep, streams=streams, sync=config.sync, + quality=config.quality, obs_keys=obs_keys, action_keys=action_keys, ): @@ -165,9 +245,24 @@ def _all_samples() -> Iterator[Sample]: produced.append(ep) episodes_done += 1 - dataset_path = Path(writer(_all_samples(), output)) + samples = _all_samples() + try: + first_sample = next(samples) + except StopIteration as error: + recorded_streams = sorted({ref.stream for ref in streams.values()}) + counts = ", ".join( + f"{stream_name}={store.stream(stream_name).count()}" + for stream_name in recorded_streams + ) + raise RuntimeError( + f"No synchronized samples were produced from {total} successful episode(s). " + f"Recorded stream counts: {counts}. Check that every configured stream records " + "data during each episode before changing the synchronization contract." + ) from error + + dataset_path = Path(selected_writer(chain((first_sample,), samples), output)) written = [e.model_copy(update={"id": f"ep_{i:06d}"}) for i, e in enumerate(produced)] - _write_dimos_meta(dataset_path, config, written) + _write_dimos_meta(dataset_path, config, written, quality_reports) logger.info( "[dataprep] succeeded — wrote %d samples across %d episodes to %s", samples_seen, @@ -179,10 +274,14 @@ def _all_samples() -> Iterator[Sample]: store.stop() -def inspect_recording(path: Path | str, status_stream: str = "status") -> dict[str, Any]: +def inspect_recording( + path: Path | str, + status_stream: str = "status", + config: DataPrepConfig | None = None, +) -> dict[str, Any]: """Summarize a source recording, including an episode left open at EOF.""" p = Path(path) - store = SqliteStore(path=str(p), must_exist=True) + store = _open_recording(p) try: stream_names = store.list_streams() stream_counts = {name: store.stream(name).count() for name in stream_names} @@ -192,7 +291,7 @@ def inspect_recording(path: Path | str, status_stream: str = "status") -> dict[s report = None episodes = report.episodes if report is not None else [] incomplete = report.incomplete if report is not None else [] - return { + result: dict[str, Any] = { "format": "recording", "path": str(p), "streams": stream_counts, @@ -202,28 +301,39 @@ def inspect_recording(path: Path | str, status_stream: str = "status") -> dict[s "discarded_episodes": sum(not episode.success for episode in episodes), "incomplete_episodes": [episode.model_dump() for episode in incomplete], } + if config is not None: + streams = {**config.observation, **config.action} + result["quality"] = [ + inspect_episode_quality( + store, episode, streams, config.sync, config.quality + ).model_dump() + for episode in episodes + if episode.success + ] + return result finally: store.stop() -def inspect_dataset(path: Path | str, fmt: str | None = None) -> dict[str, Any]: +def inspect_dataset(path: Path | str) -> dict[str, Any]: """Summarize a source recording or built dataset. Recordings report stream and episode counts plus any episode left open at EOF. Built datasets report feature shapes/dtypes, frame counts, and shape - uniformity. ``fmt`` is auto-detected when omitted. + uniformity. The format is detected from the path. """ p = Path(path) - if fmt is None: - if p.suffix == ".db": - return inspect_recording(p) - if p.suffix in (".h5", ".hdf5"): - fmt = "hdf5" - elif (p / "meta" / "info.json").exists(): - fmt = "lerobot" - else: - raise ValueError( - f"Cannot detect data format at {p}: expected a recording .db, a .hdf5 file, " - f"or a lerobot directory with meta/info.json. Pass --format explicitly." - ) - return get_inspector(fmt)(p) + if p.suffix in {".db", ".mcap"}: + return inspect_recording(p) + if p.suffix in (".h5", ".hdf5"): + from dimos.imitation.dataprep.formats.hdf5.reader import inspect + + return inspect(p) + if (p / "meta" / "info.json").exists(): + from dimos.imitation.dataprep.lerobot import inspect_lerobot_dataset + + return inspect_lerobot_dataset(p) + raise ValueError( + f"Cannot detect data format at {p}: expected a recording .db/.mcap, " + "a .hdf5 file, or a lerobot directory with meta/info.json." + ) diff --git a/dimos/imitation/dataprep/cli.py b/dimos/imitation/dataprep/cli.py deleted file mode 100644 index 9227bebb2a..0000000000 --- a/dimos/imitation/dataprep/cli.py +++ /dev/null @@ -1,110 +0,0 @@ -# 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. - -"""Implementation of the `dimos dataprep` subcommand (build + inspect). - -DataPrep is a one-shot batch transform, not a long-lived module, so it runs -as a plain command over the pure helpers in `dimos.imitation.dataprep.core` -and exits with a 0/1 status — no coordinator, no blocking loop. - -The obs/action stream maps are nested, so they come from a JSON -`DataPrepConfig` via `--config`; simple flags override `source`/`output`/ -`format` on top. See `dimos/imitation/dataprep/example_config.json`. -""" - -from __future__ import annotations - -import json -from pathlib import Path -from typing import TYPE_CHECKING, Literal - -import typer - -if TYPE_CHECKING: - from dimos.imitation.dataprep.core import DataPrepConfig - - -def _load_config( - config_path: Path | None, - source: Path | None, - output: Path | None, - output_format: Literal["lerobot", "hdf5"] | None, -) -> DataPrepConfig: - """Build a DataPrepConfig from an optional JSON file + flag overrides.""" - from dimos.imitation.dataprep.core import DataPrepConfig, OutputConfig - - if config_path is not None: - cfg = DataPrepConfig.model_validate_json(Path(config_path).read_text()) - else: - cfg = DataPrepConfig() - - updates: dict[str, object] = {} - if source is not None: - updates["source"] = str(source) - if output is not None or output_format is not None: - updates["output"] = OutputConfig( - format=output_format or cfg.output.format, - path=output or cfg.output.path, - metadata=cfg.output.metadata, - ) - return cfg.model_copy(update=updates) if updates else cfg - - -def build( - config_path: Path | None, - source: Path | None, - output: Path | None, - output_format: Literal["lerobot", "hdf5"] | None, -) -> None: - from dimos.imitation.dataprep.build import run_dataprep - - cfg = _load_config(config_path, source, output, output_format) - if not cfg.source: - typer.echo("error: no source given (use --source or set it in --config)", err=True) - raise typer.Exit(2) - if not cfg.observation and not cfg.action: - typer.echo( - "error: no observation/action streams configured; pass --config with the " - "stream maps (see dimos/imitation/dataprep/example_config.json)", - err=True, - ) - raise typer.Exit(2) - - try: - path = run_dataprep(cfg) - except Exception as e: - # CLI boundary: any failure becomes a clean message + non-zero exit - # instead of a traceback. run_dataprep raises specific errors internally. - typer.echo(f"dataprep build failed: {e}", err=True) - raise typer.Exit(1) - typer.echo(f"✓ wrote dataset to {path}") - - -def inspect(dataset: Path | None, output_format: Literal["lerobot", "hdf5"] | None) -> None: - from dimos.imitation.dataprep.build import inspect_dataset - - if dataset is None: - typer.echo( - "error: no path given (pass a recording .db, .hdf5 file, or lerobot directory)", - err=True, - ) - raise typer.Exit(2) - - try: - info = inspect_dataset(dataset, output_format) - except Exception as e: - # CLI boundary: surface failures as a message + non-zero exit, not a traceback. - typer.echo(f"dataprep inspect failed: {e}", err=True) - raise typer.Exit(1) - typer.echo(json.dumps(info, indent=2, default=str)) diff --git a/dimos/imitation/dataprep/core.py b/dimos/imitation/dataprep/core.py index 7db49cfcc1..e684270c26 100644 --- a/dimos/imitation/dataprep/core.py +++ b/dimos/imitation/dataprep/core.py @@ -14,7 +14,7 @@ """Dataset-shape types + pure helpers. -Sub-configs (StreamField, SyncConfig, OutputConfig, EpisodeExtractor) and +Sub-configs (FeatureSpec, SyncConfig, QualityConfig, OutputConfig, EpisodeExtractor) and data records (Episode, Sample) live here. So do the stateless functions that walk samples — `resolve_field`, `extract_episodes`, `iter_episode_samples`. Pure and side-effect-free; importable without @@ -27,24 +27,30 @@ from __future__ import annotations import bisect -from collections.abc import Callable, Iterator +from collections import deque +from collections.abc import Callable, Iterable, Iterator +from dataclasses import dataclass +from itertools import pairwise +import math from pathlib import Path from typing import TYPE_CHECKING, Any, Literal import numpy as np from numpy.typing import NDArray -from pydantic import BaseModel, ConfigDict, Field +from pydantic import BaseModel, ConfigDict, Field, model_validator from dimos.constants import STATE_DIR +from dimos.msgs.sensor_msgs.JointState import JointState from dimos.protocol.service.spec import BaseConfig if TYPE_CHECKING: - from dimos.memory.store.sqlite import SqliteStore + from dimos.memory.store.base import Store from dimos.memory.stream import Stream -# Each `formats//` package's writer/reader expose these, via get_writer/get_inspector. +# Each host-supported format package exposes a writer through ``get_writer``. Writer = Callable[[Iterator["Sample"], "OutputConfig"], Path] -Inspector = Callable[[Path], dict[str, Any]] + +SourceKind = Literal["snapshot", "joint_position_updates"] DEFAULT_FPS = 30.0 # resample rate == written video/timestamp rate @@ -62,19 +68,55 @@ class EpisodeExtractor(BaseConfig): ranges: list[tuple[float, float]] | None = None -class StreamField(BaseConfig): +class FeatureSpec(BaseConfig): + """Explicit dataset feature and its recorded source.""" + stream: str field: str | None = None + dtype: str + shape: tuple[int, ...] + names: list[str] + source_kind: SourceKind = "snapshot" + + @model_validator(mode="after") + def validate_schema(self) -> FeatureSpec: + if not self.shape or any(value <= 0 for value in self.shape): + raise ValueError("feature shape must contain positive dimensions") + if self.dtype == "video": + if len(self.names) != len(self.shape): + raise ValueError("video feature names must name every axis") + else: + try: + np.dtype(self.dtype) + except TypeError as error: + raise ValueError(f"unsupported feature dtype {self.dtype!r}") from error + if len(self.shape) == 1 and len(self.names) != self.shape[0]: + raise ValueError("vector feature names must match its length") + if any(not name.strip() for name in self.names): + raise ValueError("feature names must not be empty") + if self.source_kind == "joint_position_updates" and ( + self.field != "position" + or self.dtype == "video" + or len(self.shape) != 1 + or len(self.names) != len(set(self.names)) + ): + raise ValueError( + "joint_position_updates requires a position vector with unique joint names" + ) + return self class SyncConfig(BaseConfig): anchor: str - rate_hz: float - tolerance_ms: float - # TODO: add "interp" — do it per-stream (lerp low-dim vectors, force nearest - # for ndim>=3 images, can't blend frames). Only "nearest" is wired today. - strategy: Literal["nearest"] = "nearest" - action_shift: int = 1 + rate_hz: float = Field(gt=0) + tolerance_ms: float = Field(ge=0) + + +class QualityConfig(BaseConfig): + mode: Literal["strict", "fill"] = "strict" + min_source_rate_ratio: float = 0.95 + max_camera_gap_ms: float = 100.0 + max_alignment_error_ms: float = 20.0 class OutputConfig(BaseConfig): @@ -83,19 +125,34 @@ class OutputConfig(BaseConfig): metadata: dict[str, Any] = Field(default_factory=dict) -class DataPrepConfig(BaseConfig): - """Everything needed to turn a recording into a dataset. - - `source` is a recording `.db`; `observation`/`action` map dataset feature - names to recorded streams; `sync` resamples them onto a common timeline; - `output` selects format + path. Consumed by `build.run_dataprep`. - """ +class DatasetSchema(BaseConfig): + """Dataset features, episode extraction, alignment, and quality rules.""" - source: str = "" episodes: EpisodeExtractor = EpisodeExtractor() - observation: dict[str, StreamField] = Field(default_factory=dict) - action: dict[str, StreamField] = Field(default_factory=dict) + observation: dict[str, FeatureSpec] = Field(default_factory=dict) + action: dict[str, FeatureSpec] = Field(default_factory=dict) sync: SyncConfig = SyncConfig(anchor="image", rate_hz=DEFAULT_FPS, tolerance_ms=50.0) + quality: QualityConfig = QualityConfig() + + @model_validator(mode="after") + def validate_sources(self) -> DatasetSchema: + validate_source_kinds((*self.observation.values(), *self.action.values())) + return self + + +def validate_source_kinds(features: Iterable[FeatureSpec]) -> None: + """Every projection of a recorded stream must agree on its source meaning.""" + kinds: dict[str, SourceKind] = {} + for feature in features: + previous = kinds.setdefault(feature.stream, feature.source_kind) + if previous != feature.source_kind: + raise ValueError(f"stream {feature.stream!r} has conflicting source kinds") + + +class DataPrepConfig(DatasetSchema): + """Dataset interpretation plus this preparation's input and output paths.""" + + source: str = "" output: OutputConfig = OutputConfig(format="lerobot", path=STATE_DIR / "datasets" / "default") @@ -131,6 +188,20 @@ class Sample(BaseModel): observation: dict[str, NDArray[Any]] action: dict[str, NDArray[Any]] task_label: str | None = None # carried from the episode for multi-task datasets + complementary_info: dict[str, NDArray[Any]] = Field(default_factory=dict) + + +class EpisodeQualityReport(BaseModel): + episode_id: str + valid: bool + mode: Literal["strict", "fill"] + expected_frames: int = 0 + emitted_frames: int = 0 + filled_frames: int = 0 + source_rates_hz: dict[str, float] = Field(default_factory=dict) + max_gaps_ms: dict[str, float] = Field(default_factory=dict) + max_alignment_error_ms: float = 0.0 + rejection_reasons: list[str] = Field(default_factory=list) # ───────────────────────────────────────────────────────────────────────────── @@ -138,7 +209,16 @@ class Sample(BaseModel): # ───────────────────────────────────────────────────────────────────────────── -def resolve_field(msg: Any, ref: StreamField) -> NDArray[Any]: +def _joint_values(msg: JointState, field: str) -> dict[str, float]: + values = getattr(msg, field) + if len(msg.name) != len(values): + raise ValueError(f"JointState has {len(msg.name)} names but {len(values)} {field} values") + if len(msg.name) != len(set(msg.name)): + raise ValueError("JointState contains duplicate joint names") + return dict(zip(msg.name, values, strict=True)) + + +def resolve_field(msg: Any, ref: FeatureSpec) -> NDArray[Any]: """Project `msg` through `ref` (attribute access) and coerce to ndarray. Single source of truth for obs/action construction across train and @@ -148,8 +228,15 @@ def resolve_field(msg: Any, ref: StreamField) -> NDArray[Any]: - `ref.field` set: `getattr(msg, ref.field)` (or `msg[ref.field]` for dict payloads) then coerce. """ - if ref.field is None: - value: Any = msg + value: Any + if isinstance(msg, JointState) and ref.field in {"position", "velocity", "effort"}: + by_name = _joint_values(msg, ref.field) + missing = [name for name in ref.names if name not in by_name] + if missing: + raise ValueError(f"JointState is missing configured joints: {missing}") + value = [by_name[name] for name in ref.names] + elif ref.field is None: + value = msg elif isinstance(msg, dict): value = msg[ref.field] else: @@ -177,7 +264,7 @@ def is_image_array(arr: NDArray[Any]) -> bool: return False -def extract_episodes(store: SqliteStore, cfg: EpisodeExtractor) -> list[Episode]: +def extract_episodes(store: Store, cfg: EpisodeExtractor) -> list[Episode]: """Walk recorded events into Episodes per the configured strategy. EPISODE_STATUS: scan `cfg.status_stream` for state transitions emitted @@ -192,7 +279,7 @@ def extract_episodes(store: SqliteStore, cfg: EpisodeExtractor) -> list[Episode] return inspect_episodes(store, cfg).episodes -def inspect_episodes(store: SqliteStore, cfg: EpisodeExtractor) -> EpisodeReport: +def inspect_episodes(store: Store, cfg: EpisodeExtractor) -> EpisodeReport: """Extract completed episodes and retain any recording left open at EOF.""" if cfg.extractor == "ranges": if not cfg.ranges: @@ -257,11 +344,191 @@ def _commit(end_ts: float, success: bool, label: str | None) -> None: return EpisodeReport(episodes=episodes, incomplete=incomplete) +@dataclass(frozen=True) +class _AlignedFrame: + ts: float + indices: dict[str, int] + filled: bool + + +@dataclass(frozen=True) +class _AlignmentPlan: + targets: list[float] + frames: list[_AlignedFrame] + first_incomplete: float | None + max_error_s: float + + +def _alignment_plan( + features: dict[str, _FeatureSeries], + sync: SyncConfig, + quality: QualityConfig, +) -> _AlignmentPlan: + """Select source indices once for both validation and sample emission.""" + timestamps = {key: series.timestamps for key, series in features.items()} + if sync.anchor not in timestamps: + raise ValueError(f"sync.anchor {sync.anchor!r} not in streams: {sorted(timestamps)}") + + anchor = timestamps[sync.anchor] + if sync.rate_hz > 0 and anchor: + period = 1.0 / sync.rate_hz + targets: list[float] = [] + target = anchor[0] + while target <= anchor[-1]: + targets.append(target) + target += period + else: + targets = list(anchor) + + tolerance_s = min(sync.tolerance_ms, quality.max_alignment_error_ms) / 1000.0 + frames: list[_AlignedFrame] = [] + first_incomplete: float | None = None + max_error_s = 0.0 + for target in targets: + indices: dict[str, int] = {} + filled = False + for key, values in timestamps.items(): + if not values: + break + if features[key].source_kind == "joint_position_updates": + previous = bisect.bisect_right(values, target) - 1 + if previous < 0: + break + indices[key] = previous + continue + index = bisect.bisect_left(values, target) + if index == 0: + nearest = 0 + elif index == len(values): + nearest = index - 1 + else: + nearest = ( + index if values[index] - target < target - values[index - 1] else index - 1 + ) + error = abs(values[nearest] - target) + if error <= tolerance_s: + indices[key] = nearest + max_error_s = max(max_error_s, error) + continue + previous = bisect.bisect_right(values, target) - 1 + if quality.mode == "fill" and previous >= 0: + indices[key] = previous + max_error_s = max(max_error_s, abs(values[previous] - target)) + filled = True + continue + max_error_s = max(max_error_s, error) + break + if len(indices) == len(timestamps): + frames.append(_AlignedFrame(ts=target, indices=indices, filled=filled)) + elif first_incomplete is None: + first_incomplete = target + return _AlignmentPlan( + targets=targets, + frames=frames, + first_incomplete=first_incomplete, + max_error_s=max_error_s, + ) + + +@dataclass +class _FeatureSeries: + source_kind: SourceKind + timestamps: list[float] + values: list[NDArray[Any] | None] + errors: list[str | None] + + +def _source_messages( + store: Store, episode: Episode, stream: str, source_kind: SourceKind +) -> Iterator[tuple[float, Any]]: + """Read a recorded source once, independently of its feature projections.""" + start = -math.inf if source_kind == "joint_position_updates" else episode.start_ts + source: Stream[Any, Any] = store.stream(stream).time_range(start, episode.end_ts) + # Release observations as we consume them so inspection does not retain images. + pending = deque(sorted(source, key=lambda observation: observation.ts)) + while pending: + observation = pending.popleft() + yield observation.ts, observation.data + + +def _joint_position_updates( + messages: Iterator[tuple[float, Any]], start_ts: float, stream: str +) -> Iterator[tuple[float, JointState]]: + """Reconstruct effective targets; omitted joints retain their last command.""" + targets: dict[str, float] = {} + seeded = False + for timestamp, message in messages: + if not seeded and timestamp >= start_ts: + yield ( + start_ts, + JointState(ts=start_ts, name=list(targets), position=list(targets.values())), + ) + seeded = True + try: + if not isinstance(message, JointState): + raise ValueError("joint_position_updates requires JointState messages") + update = _joint_values(message, "position") + if not all(math.isfinite(value) for value in update.values()): + raise ValueError("JointState update contains non-finite values") + except (TypeError, ValueError) as error: + raise ValueError(f"{stream}: invalid update at {timestamp}: {error}") from error + targets.update(update) + if timestamp >= start_ts: + yield ( + timestamp, + JointState(ts=timestamp, name=list(targets), position=list(targets.values())), + ) + if not seeded: + yield start_ts, JointState(ts=start_ts, name=list(targets), position=list(targets.values())) + + +def _episode_features( + store: Store, + episode: Episode, + streams: dict[str, FeatureSpec], + *, + retain_values: bool, +) -> dict[str, _FeatureSeries]: + validate_source_kinds(streams.values()) + projections: dict[str, dict[str, FeatureSpec]] = {} + result = {key: _FeatureSeries(spec.source_kind, [], [], []) for key, spec in streams.items()} + for key, spec in streams.items(): + projections.setdefault(spec.stream, {})[key] = spec + for stream, specs in projections.items(): + source_kind = next(iter(specs.values())).source_kind + messages = _source_messages(store, episode, stream, source_kind) + if source_kind == "joint_position_updates": + messages = _joint_position_updates(messages, episode.start_ts, stream) + for timestamp, message in messages: + for key, spec in specs.items(): + series = result[key] + value = None + error = None + try: + value = _feature_value(message, key, spec) + except ValueError as exc: + error = str(exc) + series.timestamps.append(timestamp) + series.errors.append(error) + if retain_values: + series.values.append(value) + return result + + +def _frame_errors(frame: _AlignedFrame, features: dict[str, _FeatureSeries]) -> list[str]: + return [ + f"{error} at dataset timestamp {frame.ts}" + for key, index in frame.indices.items() + if (error := features[key].errors[index]) is not None + ] + + def iter_episode_samples( - store: SqliteStore, + store: Store, episode: Episode, - streams: dict[str, StreamField], # observation ∪ action + streams: dict[str, FeatureSpec], # observation ∪ action sync: SyncConfig, + quality: QualityConfig, obs_keys: set[str] | None = None, action_keys: set[str] | None = None, ) -> Iterator[Sample]: @@ -276,8 +543,11 @@ def iter_episode_samples( action. If omitted, every key is treated as observation (used by callers that only need raw aligned data). - With `sync.action_shift > 0` (default 1), each frame's action is taken - `action_shift` frames later (next-state target); the tail is dropped. + Held joint positions use causal command history, including before the episode. + In strict mode callers validate every target before iterating. In fill mode, + a target without a nearby value uses the last value at or before the target + and marks the frame as filled. Leading targets without a complete causal + sample are trimmed rather than fabricated. """ if sync.anchor not in streams: raise ValueError(f"sync.anchor {sync.anchor!r} not in streams: {sorted(streams)}") @@ -285,111 +555,135 @@ def iter_episode_samples( obs_keys = obs_keys if obs_keys is not None else set(streams) action_keys = action_keys if action_keys is not None else set() - tolerance_s = sync.tolerance_ms / 1000.0 - - # Materialize each stream's (timestamps, messages) once per episode. - cached: dict[str, tuple[list[float], list[Any]]] = {} - for key, ref in streams.items(): - sub: Stream[Any, Any] = store.stream(ref.stream).time_range( - episode.start_ts, episode.end_ts + features = _episode_features(store, episode, streams, retain_values=True) + plan = _alignment_plan(features, sync, quality) + for frame in plan.frames: + if _frame_errors(frame, features): + continue + obs_dict: dict[str, NDArray[Any]] = {} + act_dict: dict[str, NDArray[Any]] = {} + for key in streams: + arr = features[key].values[frame.indices[key]] + assert arr is not None + if key in action_keys: + act_dict[key] = arr + elif key in obs_keys: + obs_dict[key] = arr + yield Sample( + ts=frame.ts, + episode_id=episode.id, + observation=obs_dict, + action=act_dict, + task_label=episode.task_label, + complementary_info={"is_filled": np.asarray([frame.filled], dtype=np.bool_)}, ) - ts_list: list[float] = [] - msg_list: list[Any] = [] - for obs in sub: - ts_list.append(obs.ts) - msg_list.append(obs.data) - # Keep them sorted by time — query order is usually already sorted, but be safe. - if ts_list and any(ts_list[i] > ts_list[i + 1] for i in range(len(ts_list) - 1)): - order = sorted(range(len(ts_list)), key=ts_list.__getitem__) - ts_list = [ts_list[i] for i in order] - msg_list = [msg_list[i] for i in order] - cached[key] = (ts_list, msg_list) - - anchor_ts, _ = cached[sync.anchor] - if not anchor_ts: - return - - # Build the sequence of target timestamps for this episode. - if sync.rate_hz > 0: - # Uniform 1/rate_hz grid, phase-locked to the first anchor sample — - # what LeRobot expects (it assumes contiguous fixed-fps frames). - period = 1.0 / sync.rate_hz - targets: list[float] = [] - t = anchor_ts[0] - end = anchor_ts[-1] - while t <= end: - targets.append(t) - t += period - else: - # rate_hz=0: follow the anchor's own timestamps (no image resampling). - # dt is irregular if the camera jitters — fine for hdf5/custom trainers, - # but not LeRobot-uniform. - targets = list(anchor_ts) - - def _nearest(key: str, t: float) -> Any | None: - ts_list, msg_list = cached[key] - if not ts_list: - return None - # Nearest is i (first sample ≥ t) or i-1 (last sample < t). - i = bisect.bisect_left(ts_list, t) - if i == 0: - best = 0 - elif i == len(ts_list): - best = i - 1 - else: - best = i if (ts_list[i] - t) < (t - ts_list[i - 1]) else i - 1 - return msg_list[best] if abs(ts_list[best] - t) <= tolerance_s else None - - def _build_frames() -> Iterator[Sample]: - for t in targets: - obs_dict: dict[str, NDArray[Any]] = {} - act_dict: dict[str, NDArray[Any]] = {} - skip = False - for key, ref in streams.items(): - msg = _nearest(key, t) - if msg is None: - skip = True - break - arr = resolve_field(msg, ref) - if not is_image_array(arr): - arr = arr.astype(np.float32, copy=False) - if key in action_keys: - act_dict[key] = arr - elif key in obs_keys: - obs_dict[key] = arr - if skip: - continue - yield Sample( - ts=t, - episode_id=episode.id, - observation=obs_dict, - action=act_dict, - task_label=episode.task_label, + + +def _feature_value(msg: Any, key: str, spec: FeatureSpec) -> NDArray[Any]: + """Resolve and validate one interpreted feature for both inspect and prepare.""" + try: + value = resolve_field(msg, spec) + except (AttributeError, KeyError, TypeError, ValueError) as error: + raise ValueError(f"{key}: cannot resolve {spec.stream}.{spec.field}: {error}") from error + if tuple(value.shape) != spec.shape: + raise ValueError(f"{key}: expected shape {spec.shape}, got {tuple(value.shape)}") + if spec.dtype == "video": + if not is_image_array(value) or value.dtype != np.uint8: + raise ValueError( + f"{key}: video source must be a uint8 image, got {value.dtype} {value.shape}" ) + else: + try: + value = value.astype(np.dtype(spec.dtype), copy=False) + except (TypeError, ValueError) as error: + raise ValueError(f"{key}: cannot convert to {spec.dtype}: {error}") from error + if np.issubdtype(value.dtype, np.number) and not np.isfinite(value).all(): + raise ValueError(f"{key}: contains non-finite values") + return value - shift = max(0, sync.action_shift) - if shift == 0 or not action_keys: - yield from _build_frames() - return - # frame i keeps its obs but takes frame i+shift's action; tail dropped. - frames = list(_build_frames()) - for i in range(len(frames) - shift): - cur = frames[i] - nxt = frames[i + shift] - yield Sample( - ts=cur.ts, - episode_id=cur.episode_id, - observation=cur.observation, - action=nxt.action, - task_label=cur.task_label, +def inspect_episode_quality( + store: Store, + episode: Episode, + streams: dict[str, FeatureSpec], + sync: SyncConfig, + quality: QualityConfig, +) -> EpisodeQualityReport: + """Validate one saved episode without retaining payloads in memory.""" + report = EpisodeQualityReport(episode_id=episode.id, valid=True, mode=quality.mode) + if sync.anchor not in streams: + report.rejection_reasons.append( + f"sync.anchor {sync.anchor!r} not in features {sorted(streams)}" ) + report.valid = False + return report + + try: + features = _episode_features(store, episode, streams, retain_values=False) + except ValueError as error: + report.rejection_reasons.append(str(error)) + report.valid = False + return report + timestamps = {key: series.timestamps for key, series in features.items()} + for key, spec in streams.items(): + values = timestamps[key] + if not values: + report.rejection_reasons.append(f"{key}: stream {spec.stream!r} has no episode data") + + if spec.dtype == "video": + if len(values) > 1: + duration = values[-1] - values[0] + rate = (len(values) - 1) / duration if duration > 0 else 0.0 + gap_ms = max(b - a for a, b in pairwise(values)) * 1000.0 + else: + rate = 0.0 + gap_ms = math.inf + report.source_rates_hz[key] = rate + report.max_gaps_ms[key] = gap_ms + if quality.mode == "strict" and rate < sync.rate_hz * quality.min_source_rate_ratio: + report.rejection_reasons.append( + f"{key}: source rate {rate:.2f}Hz is below " + f"{sync.rate_hz * quality.min_source_rate_ratio:.2f}Hz" + ) + if quality.mode == "strict" and gap_ms > quality.max_camera_gap_ms: + report.rejection_reasons.append( + f"{key}: maximum source gap {gap_ms:.1f}ms exceeds " + f"{quality.max_camera_gap_ms:.1f}ms" + ) + + plan = _alignment_plan(features, sync, quality) + report.expected_frames = len(plan.targets) + reported_frame_error = False + for frame in plan.frames: + errors = _frame_errors(frame, features) + if errors: + # One failing frame explains the rejection without flooding the CLI. + if not reported_frame_error: + report.rejection_reasons.extend(errors) + reported_frame_error = True + else: + report.emitted_frames += 1 + report.filled_frames += frame.filled + report.max_alignment_error_ms = plan.max_error_s * 1000.0 + if quality.mode == "strict" and plan.first_incomplete is not None: + report.rejection_reasons.append( + f"fixed-rate target at {plan.first_incomplete:.6f} has no complete aligned sample" + ) + if not plan.targets: + report.rejection_reasons.append("anchor stream has no episode data") + if not report.emitted_frames: + report.rejection_reasons.append("episode has no complete output frames") + report.valid = not report.rejection_reasons + return report def get_writer(format_name: str) -> Writer: """Lazy-import the format writer's `write` function.""" if format_name == "lerobot": - from dimos.imitation.dataprep.formats.lerobot.writer import write + raise RuntimeError( + "LeRobot conversion requires its isolated environment; " + "use `run_lerobot_dataprep()` or `dimos imitation prepare RECORDING_DIR`" + ) elif format_name == "hdf5": from dimos.imitation.dataprep.formats.hdf5.writer import write else: @@ -397,17 +691,6 @@ def get_writer(format_name: str) -> Writer: return write -def get_inspector(format_name: str) -> Inspector: - """Lazy-import the format reader's `inspect` function.""" - if format_name == "lerobot": - from dimos.imitation.dataprep.formats.lerobot.reader import inspect - elif format_name == "hdf5": - from dimos.imitation.dataprep.formats.hdf5.reader import inspect - else: - raise ValueError(f"Unknown format: {format_name!r}") - return inspect - - def summarize_lengths(lengths: list[int]) -> dict[str, Any]: """Min/max/mean of per-episode frame counts + whether they're all equal.""" if not lengths: diff --git a/dimos/imitation/dataprep/example_config.json b/dimos/imitation/dataprep/example_config.json deleted file mode 100644 index 3458944ee0..0000000000 --- a/dimos/imitation/dataprep/example_config.json +++ /dev/null @@ -1,37 +0,0 @@ -{ - "source": "data/recordings/session.db", - "episodes": { - "extractor": "episode_status", - "status_stream": "status" - }, - "observation": { - "image": { - "stream": "color_image", - "field": "data" - }, - "joint_state": { - "stream": "coordinator_joint_state", - "field": "position" - } - }, - "action": { - "joint_target": { - "stream": "coordinator_joint_state", - "field": "position" - } - }, - "sync": { - "anchor": "image", - "rate_hz": 14.0, - "tolerance_ms": 80.0, - "action_shift": 1 - }, - "output": { - "format": "lerobot", - "path": "data/datasets/session", - "metadata": { - "robot": "xarm7", - "default_task_label": "pick_and_place" - } - } -} diff --git a/dimos/imitation/dataprep/formats/hdf5/writer.py b/dimos/imitation/dataprep/formats/hdf5/writer.py index 46557de820..d61074f728 100644 --- a/dimos/imitation/dataprep/formats/hdf5/writer.py +++ b/dimos/imitation/dataprep/formats/hdf5/writer.py @@ -26,6 +26,7 @@ timestamp (T,) float32 (T, ...) as recorded (T, ...) as recorded + complementary_info/ (T, ...) conversion metadata attrs: length, start_ts, task_index This is the ACT-original style adapted to one file with multiple episodes. @@ -67,7 +68,6 @@ def __init__(self, output: OutputConfig) -> None: self.out.parent.mkdir(parents=True, exist_ok=True) self.stats = stats_from_metadata(output.metadata) - self.default_task_label: str = output.metadata.get("default_task_label", "task") self.fps = float(output.metadata.get("fps", DEFAULT_FPS)) self.tasks_index: dict[str, int] = {} @@ -76,11 +76,12 @@ def __init__(self, output: OutputConfig) -> None: # Per-episode buffers — flushed at episode boundary. self.cur_id: str | None = None self.cur_idx = 0 - self.cur_task = self.default_task_label # actual label for the in-progress episode + self.cur_task: str | None = None self.cur_start_ts: float | None = None self.buf_ts: list[float] = [] self.buf_obs: dict[str, list[NDArray[Any]]] = {} self.buf_act: dict[str, list[NDArray[Any]]] = {} + self.buf_info: dict[str, list[NDArray[Any]]] = {} self._h5 = h5py.File(self.out, "w") self._episodes_g = self._h5.create_group("episodes") @@ -93,7 +94,9 @@ def append(self, sample: Sample) -> None: self.cur_idx += 1 self.cur_id = sample.episode_id self.cur_start_ts = float(sample.ts) - self.cur_task = sample.task_label or self.default_task_label + if not sample.task_label: + raise ValueError("every HDF5 frame requires an episode task label") + self.cur_task = sample.task_label if self.cur_task not in self.tasks_index: self.tasks_index[self.cur_task] = len(self.tasks_index) @@ -106,6 +109,8 @@ def append(self, sample: Sample) -> None: a = np.asarray(v) self.buf_act.setdefault(k, []).append(a) self.stats.update(f"action.{k}", a) + for k, v in sample.complementary_info.items(): + self.buf_info.setdefault(k, []).append(np.asarray(v)) self.total_frames += 1 def flush_episode(self) -> bool: @@ -116,6 +121,8 @@ def flush_episode(self) -> bool: ep = self._episodes_g.create_group(f"episode_{self.cur_idx:06d}") ep.attrs["length"] = len(self.buf_ts) ep.attrs["start_ts"] = float(self.cur_start_ts or 0.0) + if self.cur_task is None: + raise RuntimeError("buffered episode has no task") ep.attrs["task_index"] = self.tasks_index[self.cur_task] ep.create_dataset("timestamp", data=np.asarray(self.buf_ts, dtype=np.float32)) for k, frames in self.buf_obs.items(): @@ -131,9 +138,12 @@ def flush_episode(self) -> bool: ) for k, frames in self.buf_act.items(): ep.create_dataset(f"action/{k}", data=np.stack(frames, axis=0)) + for k, frames in self.buf_info.items(): + ep.create_dataset(f"complementary_info/{k}", data=np.stack(frames, axis=0)) self.buf_ts.clear() self.buf_obs.clear() self.buf_act.clear() + self.buf_info.clear() return True def finalize(self) -> None: diff --git a/dimos/imitation/dataprep/formats/lerobot/reader.py b/dimos/imitation/dataprep/formats/lerobot/reader.py deleted file mode 100644 index 68705991a3..0000000000 --- a/dimos/imitation/dataprep/formats/lerobot/reader.py +++ /dev/null @@ -1,87 +0,0 @@ -# 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. - -"""LeRobot v3.0 dataset reader — read-back / summary from `meta/`. - -Mirror of ``_LeRobotV3Writer``: reads ``info.json`` + the episodes parquet. -New read-side features (e.g. logging episodes to Rerun) belong here as methods. -""" - -from __future__ import annotations - -import json -from pathlib import Path -from typing import Any - -from dimos.imitation.dataprep.core import summarize_lengths -from dimos.imitation.dataprep.formats.lerobot.writer import CHUNK, EPISODES_DIR, FILE, META_DIR - -_META_COLS = {"timestamp", "frame_index", "episode_index", "index", "task_index"} - - -class _LeRobotReader: - """Read-only view over a built LeRobot v3.0 dataset. - - One instance per dataset root; lazily reads ``meta/info.json`` and the - episodes parquet on demand. - """ - - def __init__(self, path: Path) -> None: - self.root = Path(path) - self.info = json.loads((self.root / META_DIR / "info.json").read_text()) - - def _features(self) -> tuple[dict[str, Any], dict[str, Any]]: - observation: dict[str, Any] = {} - action: dict[str, Any] = {} - for name, feat in self.info.get("features", {}).items(): - if name in _META_COLS: - continue - entry = {"shape": feat.get("shape"), "dtype": feat.get("dtype")} - if name.startswith("observation"): - observation[name] = entry - elif name.startswith("action"): - action[name] = entry - return observation, action - - def _episode_lengths(self) -> list[int]: - import pyarrow.parquet as pq - - ep_file = self.root / META_DIR / EPISODES_DIR / CHUNK / f"{FILE}.parquet" - if not ep_file.exists(): - return [] - lengths: list[int] = pq.read_table(ep_file, columns=["length"]).column("length").to_pylist() - return lengths - - def summary(self) -> dict[str, Any]: - """Observation/action features (shape + dtype), episode/frame counts.""" - observation, action = self._features() - return { - "format": "lerobot", - "version": self.info.get("codebase_version"), - "path": str(self.root), - "episodes": self.info.get("total_episodes"), - "frames": self.info.get("total_frames"), - "fps": self.info.get("fps"), - "robot": self.info.get("robot_type"), - "observation": observation, - "action": action, - "episode_lengths": summarize_lengths(self._episode_lengths()), - "shapes_uniform": True, # LeRobot declares one global feature schema - "has_stats": (self.root / META_DIR / "stats.json").exists(), - } - - -def inspect(path: Path) -> dict[str, Any]: - """Summarize a LeRobot v3.0 dataset. Thin driver over ``_LeRobotReader``.""" - return _LeRobotReader(path).summary() diff --git a/dimos/imitation/dataprep/formats/lerobot/writer.py b/dimos/imitation/dataprep/formats/lerobot/writer.py deleted file mode 100644 index 4f354ad2fc..0000000000 --- a/dimos/imitation/dataprep/formats/lerobot/writer.py +++ /dev/null @@ -1,453 +0,0 @@ -# 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. - -"""LeRobot v3.0 dataset writer. - -v3.0 differs structurally from v2.x: instead of one parquet + one MP4 *per -episode*, episodes are **concatenated** into shared chunked files, and all -per-episode bookkeeping (frame/byte ranges, video time offsets, per-episode -stats) moves into an episodes *parquet*. - -Layout:: - - / - meta/info.json schema, fps, totals, features - meta/tasks.parquet task strings (indexed by `task`) - meta/stats.json aggregated per-feature stats - meta/episodes/chunk-000/file-000.parquet one row per episode (+ stats) - data/chunk-000/file-000.parquet ALL episodes' frames concatenated - videos//chunk-000/file-000.mp4 ALL episodes for a camera, concatenated - -This writer emits a **single** data file and a single MP4 per camera (chunk -000 / file 000); LeRobot supports multi-file rolling at size limits, which we -don't need yet (logged if a soft limit is exceeded). A frame's `timestamp` is -relative to its episode; the episode's `videos//from_timestamp` gives its -offset inside the shared MP4, so `from_timestamp + timestamp` locates the frame. -""" - -from __future__ import annotations - -from collections.abc import Iterator -import json -from pathlib import Path -from typing import Any - -import numpy as np -from numpy.typing import NDArray - -from dimos.imitation.dataprep.core import DEFAULT_FPS, OutputConfig, Sample, is_image_array -from dimos.imitation.dataprep.formats._stats import StreamingStats, stats_from_metadata -from dimos.utils.logging_config import setup_logger - -logger = setup_logger() - -CHUNK = "chunk-000" -FILE = "file-000" -DATA_DIR = "data" -VIDEO_DIR = "videos" -META_DIR = "meta" -EPISODES_DIR = "episodes" - -# LeRobot defaults; we write a single file but warn past these soft limits. -DATA_FILE_SIZE_MB = 100 -VIDEO_FILE_SIZE_MB = 200 -CHUNKS_SIZE = 1000 - - -def _feature_name( - prefix: str, key: str, is_image: bool, single_action: bool, single_state: bool = False -) -> str: - """Translate (prefix, key) into the LeRobot feature name. - - Canonical names lerobot policies (ACT, Diffusion, π₀) expect: - observation.state single proprio vector - action single action vector - observation.images. per-camera RGB - Multi-key fallbacks: ``observation.`` / ``action.``. - """ - if prefix == "action" and single_action: - return "action" - if is_image: - return f"observation.images.{key}" - if prefix == "observation" and single_state: - return "observation.state" - if prefix == "observation": - return f"observation.{key}" - return f"action.{key}" - - -def _nest_image_stat(vals: list[float]) -> list[list[list[float]]]: - """Per-channel [c0,c1,c2] → shape (C,1,1) [[[c0]],[[c1]],[[c2]]] (lerobot image stats).""" - return [[[float(c)]] for c in vals] - - -def _flatten_episode_stats( - final: dict[str, dict[str, Any]], feature_dtypes: dict[str, str] -) -> dict[str, Any]: - """Flatten a per-episode StreamingStats result into ``stats//`` columns. - - Image features get the (C,1,1) nesting lerobot expects; low-dim stay flat. - """ - out: dict[str, Any] = {} - for feat, entry in final.items(): - is_video = feature_dtypes.get(feat) == "video" - for k in ("mean", "std", "min", "max"): - v = entry.get(k) - if v is None: - continue - out[f"stats/{feat}/{k}"] = _nest_image_stat(v) if is_video else v - out[f"stats/{feat}/count"] = int(entry["count"]) - for q in ("q01", "q99"): - if q in entry: - out[f"stats/{feat}/{q}"] = _nest_image_stat(entry[q]) if is_video else entry[q] - return out - - -class _LeRobotV3Writer: - """Streaming writer for the LeRobot v3.0 on-disk layout. - - One instance per dataset. Drive it as ``append`` per sample, ``flush_episode`` - at each episode boundary (and once at the end), ``close`` to release the - parquet footer + MP4 handles, then ``finalize`` to emit the meta files. State - that the old single-function version threaded through ``nonlocal`` closures - lives here as instance fields, and the writer holds the lazily-imported - pyarrow/pandas/cv2 handles so the meta step needs no module-passing params. - """ - - def __init__(self, output: OutputConfig) -> None: - try: - import cv2 - except ImportError as e: - raise RuntimeError( - "LeRobot writer requires opencv-python (cv2) for MP4 encoding" - ) from e - try: - import pyarrow as pa - import pyarrow.parquet as pq - except ImportError as e: - raise RuntimeError("LeRobot writer requires pyarrow for parquet writes") from e - try: - import pandas as pd - except ImportError as e: - raise RuntimeError("LeRobot writer requires pandas for tasks.parquet") from e - - self._cv2 = cv2 - self._pa = pa - self._pq = pq - self._pd = pd - - self.output = output - self.root = Path(output.path) - (self.root / META_DIR / EPISODES_DIR / CHUNK).mkdir(parents=True, exist_ok=True) - (self.root / DATA_DIR / CHUNK).mkdir(parents=True, exist_ok=True) - - self.fps = float(output.metadata.get("fps", DEFAULT_FPS)) - self._fourcc = cv2.VideoWriter.fourcc(*"mp4v") - self.default_task_label = output.metadata.get("default_task_label", "task") - - self.global_stats = self._new_stats() # aggregated across all frames → meta/stats.json - - # Schema discovery (filled as samples flow). - self.image_keys: list[str] = [] - self.state_keys: list[str] = [] - self.action_keys: list[str] = [] - self.feature_shapes: dict[str, tuple[int, ...]] = {} - self.feature_dtypes: dict[str, str] = {} - - self.tasks_index: dict[str, int] = {} - self.episode_rows: list[dict[str, Any]] = [] - - # Single concatenated data file (opened on first flush). - self.data_path = self.root / DATA_DIR / CHUNK / f"{FILE}.parquet" - self.data_writer: Any = None - - # One MP4 per camera, persisting across episodes; from/to timestamps per episode. - self.video_writers: dict[str, Any] = {} - self.video_cum_frames: dict[str, int] = {} # frames written per camera so far - - self.global_index = 0 - self.episode_index = -1 - - # Per-episode buffers. - self.cur_id: str | None = None - self.cur_rows: list[dict[str, Any]] = [] - self.cur_ep_stats = self._new_stats() - self.cur_task = self.default_task_label # actual label for the in-progress episode - - def _new_stats(self) -> StreamingStats: - return stats_from_metadata(self.output.metadata) - - def _video_path(self, image_key: str) -> Path: - feat = _feature_name("observation", image_key, is_image=True, single_action=False) - d = self.root / VIDEO_DIR / feat / CHUNK - d.mkdir(parents=True, exist_ok=True) - return d / f"{FILE}.mp4" - - def _open_video(self, image_key: str, frame: NDArray[Any]) -> Any: - h, w = frame.shape[:2] - path = self._video_path(image_key) - vw = self._cv2.VideoWriter(str(path), self._fourcc, self.fps, (w, h)) - if not vw.isOpened(): - raise RuntimeError(f"Failed to open VideoWriter for {path}") - return vw - - def append(self, sample: Sample) -> None: - """Ingest one sample: roll over the episode if needed, update schema + - stats, append image frames to the per-camera MP4, and buffer the row.""" - cv2 = self._cv2 - if sample.episode_id != self.cur_id: - self.flush_episode() - self.cur_id = sample.episode_id - self.episode_index += 1 - self.cur_ep_stats = self._new_stats() - # Per-episode task label (falls back to the config default). - self.cur_task = sample.task_label or self.default_task_label - if self.cur_task not in self.tasks_index: - self.tasks_index[self.cur_task] = len(self.tasks_index) - - # Schema discovery + stats (global + per-episode). - n_low_dim_obs = sum( - 1 for v in sample.observation.values() if not is_image_array(np.asarray(v)) - ) - single_state = n_low_dim_obs == 1 - for k, arr in sample.observation.items(): - a = np.asarray(arr) - is_image = is_image_array(a) - name = _feature_name("observation", k, is_image, False, single_state=single_state) - if name not in self.feature_shapes: - self.feature_shapes[name] = tuple(a.shape) - self.feature_dtypes[name] = "video" if is_image else str(a.dtype) - if is_image: - if k not in self.image_keys: - self.image_keys.append(k) - elif k not in self.state_keys: - self.state_keys.append(k) - self.global_stats.update(name, a) - self.cur_ep_stats.update(name, a) - single_action = len(sample.action) == 1 - for k, arr in sample.action.items(): - a = np.asarray(arr) - name = _feature_name("action", k, is_image=False, single_action=single_action) - if name not in self.feature_shapes: - self.feature_shapes[name] = tuple(a.shape) - self.feature_dtypes[name] = str(a.dtype) - if k not in self.action_keys: - self.action_keys.append(k) - self.global_stats.update(name, a) - self.cur_ep_stats.update(name, a) - - # Append image frames to the per-camera MP4 (RGB→BGR; cv2 is BGR-native). - for k, arr in sample.observation.items(): - a = np.asarray(arr) - if is_image_array(a): - if k not in self.video_writers: - self.video_writers[k] = self._open_video(k, a) - if a.ndim == 2: # grayscale → 3-channel BGR for the MP4 - bgr = cv2.cvtColor(a, cv2.COLOR_GRAY2BGR) - elif a.shape[-1] == 3: # RGB → BGR (cv2 is BGR-native) - bgr = cv2.cvtColor(a, cv2.COLOR_RGB2BGR) - else: - bgr = a - self.video_writers[k].write(bgr) - self.video_cum_frames[k] = self.video_cum_frames.get(k, 0) + 1 - - frame_index = len(self.cur_rows) - self.cur_rows.append( - { - "timestamp": frame_index / self.fps, # relative to this episode - "frame_index": frame_index, - "episode_index": self.episode_index, - "index": self.global_index, - "task_index": self.tasks_index[self.cur_task], - "obs": { - k: np.asarray(v) - for k, v in sample.observation.items() - if not is_image_array(np.asarray(v)) - }, - "act": {k: np.asarray(v) for k, v in sample.action.items()}, - } - ) - self.global_index += 1 - - def flush_episode(self) -> None: - """Write the buffered episode's rows to the concatenated data parquet and - append its metadata row. No-op when the buffer is empty.""" - if not self.cur_rows: - return - pa = self._pa - cur_rows = self.cur_rows - length = len(cur_rows) - single_state = len(self.state_keys) == 1 - single_action = len(self.action_keys) == 1 - - cols: dict[str, Any] = { - "timestamp": pa.array([r["timestamp"] for r in cur_rows], pa.float32()), - "frame_index": pa.array([r["frame_index"] for r in cur_rows], pa.int64()), - "episode_index": pa.array([r["episode_index"] for r in cur_rows], pa.int64()), - "index": pa.array([r["index"] for r in cur_rows], pa.int64()), - "task_index": pa.array([r["task_index"] for r in cur_rows], pa.int64()), - } - f32_list = pa.list_(pa.float32()) - for k in self.state_keys: - name = _feature_name("observation", k, False, False, single_state=single_state) - cols[name] = pa.array([r["obs"][k].tolist() for r in cur_rows], type=f32_list) - for k in self.action_keys: - name = _feature_name("action", k, False, single_action=single_action) - cols[name] = pa.array([r["act"][k].tolist() for r in cur_rows], type=f32_list) - table = pa.Table.from_pydict(cols) - if self.data_writer is None: - self.data_writer = self._pq.ParquetWriter( - self.data_path, table.schema, compression="snappy" - ) - self.data_writer.write_table(table) - - # Episode metadata row. - row: dict[str, Any] = { - "episode_index": self.episode_index, - "tasks": [list(self.tasks_index.keys())[cur_rows[0]["task_index"]]], - "length": length, - "data/chunk_index": 0, - "data/file_index": 0, - "dataset_from_index": self.global_index - length, - "dataset_to_index": self.global_index, - "meta/episodes/chunk_index": 0, - "meta/episodes/file_index": 0, - } - for k in self.image_keys: - feat = _feature_name("observation", k, is_image=True, single_action=False) - cum = self.video_cum_frames.get(k, 0) - row[f"videos/{feat}/chunk_index"] = 0 - row[f"videos/{feat}/file_index"] = 0 - row[f"videos/{feat}/from_timestamp"] = (cum - length) / self.fps - row[f"videos/{feat}/to_timestamp"] = cum / self.fps - row.update(_flatten_episode_stats(self.cur_ep_stats.finalize(), self.feature_dtypes)) - self.episode_rows.append(row) - cur_rows.clear() - - def close(self) -> None: - """Release the parquet footer and MP4 handles. Safe to call on partial - writes — without this the data file has no footer and is unreadable.""" - if self.data_writer is not None: - self.data_writer.close() - self.data_writer = None - for vw in self.video_writers.values(): - vw.release() - self.video_writers.clear() - - def finalize(self) -> None: - """Write info.json, tasks.parquet, episodes parquet, and aggregated stats.json.""" - pa, pq, pd = self._pa, self._pq, self._pd - total_episodes = len(self.episode_rows) - total_frames = self.global_index - if self.data_path.exists() and self.data_path.stat().st_size > DATA_FILE_SIZE_MB * 1e6: - logger.warning( - "[dataprep] data file exceeds %d MB (single-file writer, no rolling): %s", - DATA_FILE_SIZE_MB, - self.data_path, - ) - - features: dict[str, Any] = {} - for name, shape in self.feature_shapes.items(): - if self.feature_dtypes[name] == "video": - features[name] = { - "dtype": "video", - "shape": list(shape), - "names": ["height", "width", "channel"], - "info": { - "video.fps": self.fps, - "video.height": int(shape[0]), - "video.width": int(shape[1]), - "video.channels": int(shape[2]) if len(shape) > 2 else 3, - "video.codec": "mp4v", - "video.pix_fmt": "yuv420p", - "video.is_depth_map": False, - "has_audio": False, - }, - } - else: - n = int(shape[0]) if shape else 0 - base = name.split(".")[-1] - features[name] = { - "dtype": self.feature_dtypes[name], - "shape": list(shape), - "names": [f"{base}_{i}" for i in range(n)], - } - for col, dt in [ - ("timestamp", "float32"), - ("frame_index", "int64"), - ("episode_index", "int64"), - ("index", "int64"), - ("task_index", "int64"), - ]: - features[col] = {"dtype": dt, "shape": [1], "names": None} - - info = { - "codebase_version": "v3.0", - "robot_type": self.output.metadata.get("robot", "unknown"), - "total_episodes": total_episodes, - "total_frames": total_frames, - "total_tasks": len(self.tasks_index), - "chunks_size": CHUNKS_SIZE, - "data_files_size_in_mb": DATA_FILE_SIZE_MB, - "video_files_size_in_mb": VIDEO_FILE_SIZE_MB, - "fps": self.fps, - "splits": {"train": f"0:{total_episodes}"}, - "data_path": "data/chunk-{chunk_index:03d}/file-{file_index:03d}.parquet", - "video_path": "videos/{video_key}/chunk-{chunk_index:03d}/file-{file_index:03d}.mp4", - "features": features, - } - with open(self.root / META_DIR / "info.json", "w") as f: - json.dump(info, f, indent=2) - - # tasks.parquet — task strings as the (named) index + a task_index column. - tasks_df = pd.DataFrame( - {"task_index": list(self.tasks_index.values())}, - index=pd.Index(list(self.tasks_index.keys()), name="task"), - ) - tasks_df.to_parquet(self.root / META_DIR / "tasks.parquet") - - # episodes parquet — one row per episode (+ flattened per-episode stats). - ep_table = pa.Table.from_pylist(self.episode_rows) - pq.write_table( - ep_table, - self.root / META_DIR / EPISODES_DIR / CHUNK / f"{FILE}.parquet", - compression="snappy", - ) - - # Aggregated stats.json (image features nested to (C,1,1)). - final_stats = self.global_stats.finalize() - for name, entry in final_stats.items(): - if self.feature_dtypes.get(name) == "video": - for k in ("mean", "std", "min", "max"): - if entry.get(k) is not None: - entry[k] = _nest_image_stat(entry[k]) - with open(self.root / META_DIR / "stats.json", "w") as f: - json.dump(final_stats, f, indent=2) - - -def write(samples: Iterator[Sample], output: OutputConfig) -> Path: - """Drain `samples`, write a LeRobot v3.0 dataset. Returns the dataset root path.""" - writer = _LeRobotV3Writer(output) - # try/finally so the parquet footer is written and MP4s are released even if - # the drain raises mid-stream — otherwise the data file is unreadable (no - # footer) and the videos lose their index. - try: - for sample in samples: - writer.append(sample) - writer.flush_episode() - finally: - writer.close() - - writer.finalize() - return writer.root diff --git a/dimos/imitation/dataprep/formats/test_hdf5.py b/dimos/imitation/dataprep/formats/test_hdf5.py index 7341344b29..7ef6caa54c 100644 --- a/dimos/imitation/dataprep/formats/test_hdf5.py +++ b/dimos/imitation/dataprep/formats/test_hdf5.py @@ -42,6 +42,8 @@ def _samples(n_episodes: int = 2, n_frames: int = 3) -> Iterator[Sample]: episode_id=f"ep_{ep:06d}", observation={"state": (np.arange(4, dtype=np.float32) + i)}, action={"action": np.full(2, float(i), dtype=np.float32)}, + task_label=f"task {ep}", + complementary_info={"is_filled": np.asarray([i == 1], dtype=np.bool_)}, ) @@ -66,6 +68,11 @@ def test_hdf5_roundtrip_counts_and_shapes(tmp_path: Path) -> None: assert info["shapes_uniform"] is True assert info["has_stats"] is True assert info["episode_lengths"] == {"min": 3, "max": 3, "mean": 3.0, "uniform": True} + with h5py.File(path, "r") as dataset: + np.testing.assert_array_equal( + dataset["episodes/episode_000000/complementary_info/is_filled"][:], + [[False], [True], [False]], + ) def test_hdf5_extension_appended_when_missing(tmp_path: Path) -> None: diff --git a/dimos/imitation/dataprep/formats/test_lerobot.py b/dimos/imitation/dataprep/formats/test_lerobot.py deleted file mode 100644 index 1a26b6ade0..0000000000 --- a/dimos/imitation/dataprep/formats/test_lerobot.py +++ /dev/null @@ -1,220 +0,0 @@ -# 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. - -"""Smoke tests for the LeRobot v3.0 writer/reader. - -Asserts the v3.0 layout: a single concatenated data parquet, parquet meta -(tasks + episodes, no jsonl), and one MP4 per camera under -`videos//chunk-000/`. pyarrow/pandas (the `learning` extra) and cv2 are -test dependencies, so these always run. -""" - -from __future__ import annotations - -from collections.abc import Iterator -import json -from pathlib import Path - -import numpy as np -import pandas as pd -import pyarrow.parquet as pq -import pytest - -from dimos.imitation.dataprep.core import OutputConfig, Sample -from dimos.imitation.dataprep.formats.lerobot.reader import inspect -from dimos.imitation.dataprep.formats.lerobot.writer import write - - -def _state_samples(n: int = 4) -> Iterator[Sample]: - for i in range(n): - yield Sample( - ts=float(i), - episode_id="ep_000000", - observation={"state": np.arange(6, dtype=np.float32)}, - action={"action": np.full(6, float(i), dtype=np.float32)}, - ) - - -def _two_episode_samples() -> Iterator[Sample]: - for ep in range(2): - for i in range(3): - yield Sample( - ts=float(ep * 3 + i), - episode_id=f"ep_{ep:06d}", - observation={"state": np.arange(6, dtype=np.float32) + ep}, - action={"action": np.full(6, float(i), dtype=np.float32)}, - ) - - -def _image_samples(n: int = 4) -> Iterator[Sample]: - for i in range(n): - yield Sample( - ts=float(i), - episode_id="ep_000000", - observation={ - "state": np.arange(6, dtype=np.float32), - "cam": np.full((16, 16, 3), i, dtype=np.uint8), - }, - action={"action": np.zeros(6, dtype=np.float32)}, - ) - - -def test_lerobot_v3_state_only_layout_and_naming(tmp_path: Path) -> None: - out = OutputConfig( - format="lerobot", path=tmp_path / "ds", metadata={"fps": 10.0, "robot": "xarm7"} - ) - root = write(_state_samples(), out) - - # v3.0: concatenated single data file + parquet meta (no jsonl, no per-episode parquet) - assert (root / "meta" / "info.json").exists() - assert (root / "meta" / "tasks.parquet").exists() - assert (root / "meta" / "stats.json").exists() - assert (root / "meta" / "episodes" / "chunk-000" / "file-000.parquet").exists() - assert (root / "data" / "chunk-000" / "file-000.parquet").exists() - assert not (root / "meta" / "episodes.jsonl").exists() - assert not (root / "meta" / "tasks.jsonl").exists() - - info = json.loads((root / "meta" / "info.json").read_text()) - assert info["codebase_version"] == "v3.0" - assert info["total_episodes"] == 1 - assert info["total_frames"] == 4 - assert info["fps"] == 10.0 - assert info["data_path"] == "data/chunk-{chunk_index:03d}/file-{file_index:03d}.parquet" - # single low-dim state + single action → canonical names - assert "observation.state" in info["features"] - assert "action" in info["features"] - - -def test_lerobot_v3_episode_metadata_columns(tmp_path: Path) -> None: - out = OutputConfig(format="lerobot", path=tmp_path / "ds", metadata={"fps": 10.0}) - # two episodes so dataset_from/to_index advance - root = write(_two_episode_samples(), out) - ep = pq.read_table(root / "meta" / "episodes" / "chunk-000" / "file-000.parquet") - cols = set(ep.column_names) - for required in ( - "episode_index", - "tasks", - "length", - "dataset_from_index", - "dataset_to_index", - "data/chunk_index", - "data/file_index", - "meta/episodes/chunk_index", - "meta/episodes/file_index", - ): - assert required in cols, f"missing episode column {required}" - # per-episode stats are embedded (flattened) - assert any(c.startswith("stats/observation.state/") for c in cols) - rows = ep.to_pylist() - assert [r["episode_index"] for r in rows] == [0, 1] - assert rows[0]["dataset_from_index"] == 0 and rows[0]["dataset_to_index"] == 3 - assert rows[1]["dataset_from_index"] == 3 and rows[1]["dataset_to_index"] == 6 - - -def test_lerobot_v3_writer_closed_on_midstream_error(tmp_path: Path) -> None: - """If the drain raises after an episode was flushed, the data parquet must - still be readable (footer written by the finally), not a headerless stub.""" - - def bad_samples() -> Iterator[Sample]: - for i in range(3): # episode 0 - yield Sample( - ts=float(i), - episode_id="ep_000000", - observation={"state": np.arange(6, dtype=np.float32)}, - action={"action": np.full(6, float(i), dtype=np.float32)}, - ) - # first frame of episode 1 flushes episode 0 (opens + writes the parquet)… - yield Sample( - ts=3.0, - episode_id="ep_000001", - observation={"state": np.arange(6, dtype=np.float32)}, - action={"action": np.zeros(6, dtype=np.float32)}, - ) - raise RuntimeError("boom mid-stream") # …then blow up before the final flush - - out = OutputConfig(format="lerobot", path=tmp_path / "ds", metadata={"fps": 10.0}) - with pytest.raises(RuntimeError, match="boom"): - write(bad_samples(), out) - - # episode 0's 3 frames were flushed; the file must have a valid footer. - data = tmp_path / "ds" / "data" / "chunk-000" / "file-000.parquet" - assert data.exists() - assert pq.read_table(data).num_rows == 3 # raises ArrowInvalid if footer missing - - -def test_lerobot_v3_per_episode_task_labels(tmp_path: Path) -> None: - """Episodes with distinct task_labels must produce distinct tasks + task_index - (multi-task recordings must not collapse to one task).""" - - def samples() -> Iterator[Sample]: - for ep, task in ((0, "pick"), (1, "place")): - for i in range(3): - yield Sample( - ts=float(ep * 3 + i), - episode_id=f"ep_{ep:06d}", - observation={"state": np.arange(6, dtype=np.float32)}, - action={"action": np.zeros(6, dtype=np.float32)}, - task_label=task, - ) - - out = OutputConfig(format="lerobot", path=tmp_path / "ds", metadata={"fps": 10.0}) - root = write(samples(), out) - - tasks = pd.read_parquet(root / "meta" / "tasks.parquet") - assert set(tasks.index) == {"pick", "place"} - - ep = pq.read_table(root / "meta" / "episodes" / "chunk-000" / "file-000.parquet").to_pylist() - assert ep[0]["tasks"] == ["pick"] - assert ep[1]["tasks"] == ["place"] - - data = pq.read_table(root / "data" / "chunk-000" / "file-000.parquet") - ti = data.column("task_index").to_pylist() - assert ti[:3] == [0, 0, 0] # episode 0 → task 0 (pick) - assert ti[3:] == [1, 1, 1] # episode 1 → task 1 (place) - - -def test_lerobot_v3_inspect_state_only(tmp_path: Path) -> None: - out = OutputConfig(format="lerobot", path=tmp_path / "ds", metadata={"fps": 10.0}) - root = write(_state_samples(), out) - info = inspect(root) - assert info["format"] == "lerobot" - assert info["version"] == "v3.0" - assert info["episodes"] == 1 - assert info["frames"] == 4 - assert "observation.state" in info["observation"] - assert "action" in info["action"] - assert info["has_stats"] is True - - -def test_lerobot_v3_with_images_writes_concatenated_mp4(tmp_path: Path) -> None: - out = OutputConfig(format="lerobot", path=tmp_path / "ds", metadata={"fps": 10.0}) - try: - root = write(_image_samples(), out) - except RuntimeError as e: - if "VideoWriter" in str(e): - pytest.skip(f"no mp4v encoder available in this environment: {e}") - raise - - # v3.0 video path: videos//chunk-000/file-000.mp4 (key before chunk, one per camera) - mp4 = root / "videos" / "observation.images.cam" / "chunk-000" / "file-000.mp4" - assert mp4.exists() and mp4.stat().st_size > 0 - - info = json.loads((root / "meta" / "info.json").read_text()) - assert ( - info["video_path"] == "videos/{video_key}/chunk-{chunk_index:03d}/file-{file_index:03d}.mp4" - ) - assert info["features"]["observation.images.cam"]["dtype"] == "video" - # image column is excluded from parquet; state/action remain - assert "observation.state" in info["features"] - assert info["total_frames"] == 4 diff --git a/dimos/imitation/dataprep/lerobot.py b/dimos/imitation/dataprep/lerobot.py new file mode 100644 index 0000000000..f73431ad4c --- /dev/null +++ b/dimos/imitation/dataprep/lerobot.py @@ -0,0 +1,101 @@ +# 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. + +"""Launch native LeRobot conversion in the policy runtime environment.""" + +from __future__ import annotations + +from pathlib import Path +import subprocess +from typing import Any + +from dimos.experimental.isolated_python.module import ( + isolated_python_environment, + isolated_python_run_command, +) +from dimos.imitation.dataprep._lerobot_protocol import ( + RESULT_ADAPTER, + BuildRequest, + BuildResult, + InspectRequest, + InspectResult, + Request, + Result, +) +from dimos.imitation.dataprep.core import DataPrepConfig +from dimos.imitation.policy.lerobot.module import LeRobotPolicyModule +from dimos.utils.cache import cache_usage_guard +from dimos.utils.data import get_project_root + + +def lerobot_project() -> Path: + """Locate the LeRobot project in the shared checkout.""" + return get_project_root() / LeRobotPolicyModule.project_dir + + +def _run(request: Request) -> Result: + """Run one typed request under the locked LeRobot dependency stack.""" + project = lerobot_project() + command = isolated_python_run_command( + project, + "python", + "-m", + "dimos_lerobot.dataprep", + ) + env = isolated_python_environment(project) + try: + with cache_usage_guard(): + result = subprocess.run( + command, + cwd=project, + env=env, + input=request.model_dump_json(), + capture_output=True, + text=True, + ) + except FileNotFoundError as error: + raise RuntimeError( + "uv is required for LeRobot dataprep; install uv and ensure it is on PATH" + ) from error + if result.returncode: + output = result.stderr.strip() or result.stdout.strip() + raise RuntimeError(f"LeRobot dataprep exited with status {result.returncode}: {output}") + try: + return RESULT_ADAPTER.validate_json(result.stdout.rstrip().rsplit("\n", 1)[-1]) + except ValueError as error: + raise RuntimeError( + f"LeRobot dataprep returned an invalid result: {result.stdout!r}" + ) from error + + +def run_lerobot_dataprep(config: DataPrepConfig) -> Path: + """Build a dataset in the isolated LeRobot environment.""" + config = config.model_copy( + update={ + "source": str(Path(config.source).resolve()), + "output": config.output.model_copy(update={"path": config.output.path.resolve()}), + } + ) + result = _run(BuildRequest(config=config)) + if not isinstance(result, BuildResult): + raise RuntimeError(f"LeRobot dataprep returned {result.command!r} for a build request") + return result.path + + +def inspect_lerobot_dataset(path: Path) -> dict[str, Any]: + """Inspect a dataset in the isolated LeRobot environment.""" + result = _run(InspectRequest(path=path.resolve())) + if not isinstance(result, InspectResult): + raise RuntimeError(f"LeRobot dataprep returned {result.command!r} for an inspect request") + return result.info diff --git a/dimos/imitation/dataprep/test_core.py b/dimos/imitation/dataprep/test_core.py index 5e9492b73b..e5a7ab3436 100644 --- a/dimos/imitation/dataprep/test_core.py +++ b/dimos/imitation/dataprep/test_core.py @@ -32,12 +32,17 @@ from dimos.imitation.dataprep.build import _write_dimos_meta, inspect_dataset, run_dataprep from dimos.imitation.dataprep.core import ( DataPrepConfig, + DatasetSchema, Episode, EpisodeExtractor, + EpisodeQualityReport, + FeatureSpec, OutputConfig, - StreamField, + QualityConfig, + Sample, SyncConfig, extract_episodes, + inspect_episode_quality, inspect_episodes, is_image_array, iter_episode_samples, @@ -45,6 +50,7 @@ summarize_lengths, ) from dimos.memory.store.sqlite import SqliteStore +from dimos.msgs.sensor_msgs.JointState import JointState @pytest.mark.parametrize( @@ -126,6 +132,10 @@ def _status(events: list[tuple[float, str, str | None]]) -> list[_Obs]: return [_Obs(ts=ts, data=_Status(last_event=ev, task_label=lbl)) for ts, ev, lbl in events] +def _feature(stream: str, field: str | None = "position") -> FeatureSpec: + return FeatureSpec(stream=stream, field=field, dtype="float32", shape=(1,), names=["joint"]) + + # ── resolve_field ──────────────────────────────────────────────────────────── @@ -134,19 +144,93 @@ def test_resolve_field_attribute() -> None: class Msg: position: list[float] - arr = resolve_field(Msg(position=[1.0, 2.0, 3.0]), StreamField(stream="x", field="position")) + arr = resolve_field( + Msg(position=[1.0, 2.0, 3.0]), + FeatureSpec( + stream="x", field="position", dtype="float32", shape=(3,), names=["a", "b", "c"] + ), + ) assert isinstance(arr, np.ndarray) np.testing.assert_array_equal(arr, np.array([1.0, 2.0, 3.0])) +def test_resolve_field_orders_joint_state_by_feature_names() -> None: + arr = resolve_field( + JointState(name=["joint_b", "joint_a"], position=[2.0, 1.0]), + FeatureSpec( + stream="state", + field="position", + dtype="float32", + shape=(2,), + names=["joint_a", "joint_b"], + ), + ) + + np.testing.assert_array_equal(arr, np.array([1.0, 2.0])) + + +@pytest.mark.parametrize( + ("field", "expected"), + [("velocity", [1.0, 2.0]), ("effort", [3.0, 4.0])], +) +def test_resolve_field_orders_all_joint_state_vectors(field: str, expected: list[float]) -> None: + state = JointState( + name=["joint_b", "joint_a"], + velocity=[2.0, 1.0], + effort=[4.0, 3.0], + ) + spec = FeatureSpec( + stream="state", + field=field, + dtype="float32", + shape=(2,), + names=["joint_a", "joint_b"], + ) + + np.testing.assert_array_equal(resolve_field(state, spec), expected) + + +def test_resolve_field_rejects_duplicate_joint_names() -> None: + with pytest.raises(ValueError, match="duplicate joint names"): + resolve_field( + JointState(name=["joint_a", "joint_a"], position=[1.0, 2.0]), + FeatureSpec( + stream="state", + field="position", + dtype="float32", + shape=(1,), + names=["joint_a"], + ), + ) + + +def test_resolve_field_rejects_missing_configured_joint() -> None: + with pytest.raises(ValueError, match="missing configured joints.*joint_b"): + resolve_field( + JointState(name=["joint_a"], position=[1.0]), + FeatureSpec( + stream="state", + field="position", + dtype="float32", + shape=(2,), + names=["joint_a", "joint_b"], + ), + ) + + def test_resolve_field_dict_payload() -> None: - arr = resolve_field({"q": [4, 5]}, StreamField(stream="x", field="q")) + arr = resolve_field( + {"q": [4, 5]}, + FeatureSpec(stream="x", field="q", dtype="float32", shape=(2,), names=["a", "b"]), + ) np.testing.assert_array_equal(arr, np.array([4, 5])) def test_resolve_field_none_passthrough_ndarray() -> None: src = np.arange(6).reshape(2, 3) - out = resolve_field(src, StreamField(stream="x", field=None)) + out = resolve_field( + src, FeatureSpec(stream="x", dtype="float32", shape=(2, 3), names=["row", "column"]) + ) assert out is src # ndarray passes straight through @@ -156,7 +240,9 @@ class Image: data: np.ndarray img = Image(data=np.ones((2, 2))) - out = resolve_field(img, StreamField(stream="x", field=None)) + out = resolve_field( + img, FeatureSpec(stream="x", dtype="float32", shape=(2, 2), names=["row", "column"]) + ) np.testing.assert_array_equal(out, np.ones((2, 2))) @@ -255,8 +341,7 @@ class S: return [_Obs(ts=ts, data=S(position=[v])) for ts, v in values] -def test_sync_basic_no_shift() -> None: - # obs == action, shift disabled → one sample per anchor target +def test_sync_uses_same_frame_observation_and_applied_action() -> None: store = _FakeStore( { "js": _scalar_stream([(0.0, 10.0), (1.0, 11.0), (2.0, 12.0)]), @@ -264,71 +349,431 @@ def test_sync_basic_no_shift() -> None: ) ep = Episode(id="ep_0", start_ts=0.0, end_ts=2.0) streams = { - "state": StreamField(stream="js", field="position"), - "act": StreamField(stream="js", field="position"), + "state": _feature("js"), + "act": _feature("js"), } - sync = SyncConfig(anchor="state", rate_hz=1.0, tolerance_ms=100.0, action_shift=0) + sync = SyncConfig(anchor="state", rate_hz=1.0, tolerance_ms=100.0) samples = list( - iter_episode_samples(store, ep, streams, sync, obs_keys={"state"}, action_keys={"act"}) + iter_episode_samples( + store, + ep, + streams, + sync, + QualityConfig(), + obs_keys={"state"}, + action_keys={"act"}, + ) ) assert len(samples) == 3 # action equals state at the same frame np.testing.assert_array_equal(samples[0].observation["state"], samples[0].action["act"]) -def test_sync_action_shift_next_state() -> None: - store = _FakeStore({"js": _scalar_stream([(0.0, 10.0), (1.0, 11.0), (2.0, 12.0)])}) +def test_fill_preserves_grid_and_marks_held_frame() -> None: + # anchor ticks every 1s, but the second stream has a big gap around t=1 + store = _FakeStore( + { + "anchor": _scalar_stream([(0.0, 0.0), (1.0, 0.0), (2.0, 0.0)]), + "other": _scalar_stream([(0.0, 5.0), (2.0, 7.0)]), # nothing near t=1 + } + ) ep = Episode(id="ep_0", start_ts=0.0, end_ts=2.0) streams = { - "state": StreamField(stream="js", field="position"), - "act": StreamField(stream="js", field="position"), + "anchor": _feature("anchor"), + "other": _feature("other"), } - sync = SyncConfig(anchor="state", rate_hz=1.0, tolerance_ms=100.0, action_shift=1) + sync = SyncConfig(anchor="anchor", rate_hz=1.0, tolerance_ms=100.0) samples = list( - iter_episode_samples(store, ep, streams, sync, obs_keys={"state"}, action_keys={"act"}) + iter_episode_samples( + store, + ep, + streams, + sync, + QualityConfig(mode="fill"), + obs_keys={"anchor", "other"}, + ) ) - # 3 frames, shift 1 → 2 emitted; trailing frame dropped - assert len(samples) == 2 - # frame 0: obs is state@0 (10), action is state@1 (11) - np.testing.assert_array_equal(samples[0].observation["state"], [10.0]) - np.testing.assert_array_equal(samples[0].action["act"], [11.0]) - np.testing.assert_array_equal(samples[1].observation["state"], [11.0]) - np.testing.assert_array_equal(samples[1].action["act"], [12.0]) + assert [round(s.ts) for s in samples] == [0, 1, 2] + assert [bool(s.complementary_info["is_filled"][0]) for s in samples] == [False, True, False] + np.testing.assert_array_equal(samples[1].observation["other"], [5.0]) -def test_sync_tolerance_skips_unmatched_frame() -> None: - # anchor ticks every 1s, but the second stream has a big gap around t=1 +def test_strict_quality_rejects_missing_fixed_rate_slot() -> None: store = _FakeStore( { - "anchor": _scalar_stream([(0.0, 0.0), (1.0, 0.0), (2.0, 0.0)]), - "other": _scalar_stream([(0.0, 5.0), (2.0, 7.0)]), # nothing near t=1 + "anchor": _scalar_stream([(0.0, 0.0), (1.0, 1.0), (2.0, 2.0)]), + "other": _scalar_stream([(0.0, 5.0), (2.0, 7.0)]), } ) - ep = Episode(id="ep_0", start_ts=0.0, end_ts=2.0) - streams = { - "anchor": StreamField(stream="anchor", field="position"), - "other": StreamField(stream="other", field="position"), - } - sync = SyncConfig(anchor="anchor", rate_hz=1.0, tolerance_ms=100.0, action_shift=0) - samples = list(iter_episode_samples(store, ep, streams, sync, obs_keys={"anchor", "other"})) - # t=1 dropped (no `other` within 100ms) → only t=0 and t=2 survive - assert [round(s.ts) for s in samples] == [0, 2] + episode = Episode(id="ep_0", start_ts=0.0, end_ts=2.0) + streams = {"anchor": _feature("anchor"), "other": _feature("other")} + + report = inspect_episode_quality( + store, + episode, + streams, + SyncConfig(anchor="anchor", rate_hz=1.0, tolerance_ms=20.0), + QualityConfig(mode="strict"), + ) + + assert report.valid is False + assert report.expected_frames == 3 + assert report.emitted_frames == 2 + assert "no complete aligned sample" in report.rejection_reasons[-1] + + samples = list( + iter_episode_samples( + store, + episode, + streams, + SyncConfig(anchor="anchor", rate_hz=1.0, tolerance_ms=20.0), + QualityConfig(mode="strict"), + ) + ) + assert len(samples) == report.emitted_frames + + +def test_fill_quality_accepts_gap_and_reports_filled_slot() -> None: + store = _FakeStore( + { + "anchor": _scalar_stream([(0.0, 0.0), (1.0, 1.0), (2.0, 2.0)]), + "other": _scalar_stream([(0.0, 5.0), (2.0, 7.0)]), + } + ) + episode = Episode(id="ep_0", start_ts=0.0, end_ts=2.0) + streams = {"anchor": _feature("anchor"), "other": _feature("other")} + + report = inspect_episode_quality( + store, + episode, + streams, + SyncConfig(anchor="anchor", rate_hz=1.0, tolerance_ms=20.0), + QualityConfig(mode="fill"), + ) + + assert report.valid is True + assert report.expected_frames == 3 + assert report.emitted_frames == 3 + assert report.filled_frames == 1 + + samples = list( + iter_episode_samples( + store, + episode, + streams, + SyncConfig(anchor="anchor", rate_hz=1.0, tolerance_ms=20.0), + QualityConfig(mode="fill"), + ) + ) + assert len(samples) == report.emitted_frames + assert sum(bool(sample.complementary_info["is_filled"][0]) for sample in samples) == ( + report.filled_frames + ) def test_sync_missing_anchor_raises() -> None: ep = Episode(id="ep_0", start_ts=0.0, end_ts=1.0) - streams = {"x": StreamField(stream="x", field="position")} + streams = {"x": _feature("x")} sync = SyncConfig(anchor="not_there", rate_hz=1.0, tolerance_ms=10.0) with pytest.raises(ValueError, match="anchor"): - list(iter_episode_samples(_FakeStore({}), ep, streams, sync)) + list(iter_episode_samples(_FakeStore({}), ep, streams, sync, QualityConfig())) def test_sync_empty_anchor_yields_nothing() -> None: store = _FakeStore({"a": []}) ep = Episode(id="ep_0", start_ts=0.0, end_ts=1.0) - streams = {"a": StreamField(stream="a", field="position")} + streams = {"a": _feature("a")} sync = SyncConfig(anchor="a", rate_hz=1.0, tolerance_ms=10.0) - assert list(iter_episode_samples(store, ep, streams, sync)) == [] + assert list(iter_episode_samples(store, ep, streams, sync, QualityConfig())) == [] + + +def _held_action(): + return FeatureSpec( + stream="commands", + field="position", + dtype="float32", + shape=(3,), + names=["left", "right", "gripper"], + source_kind="joint_position_updates", + ) + + +@pytest.mark.parametrize("sampling", ["nearest", "joint_position_hold"]) +def test_obsolete_sampling_option_is_rejected(sampling): + values = _held_action().model_dump() + values["sampling"] = sampling + with pytest.raises( + ValueError, + match="Extra inputs are not permitted", + ): + FeatureSpec(**values) + + +def test_saved_schema_rejects_conflicting_source_kinds(): + updates = _held_action() + snapshot = FeatureSpec(**(updates.model_dump() | {"source_kind": "snapshot"})) + with pytest.raises(ValueError, match="conflicting source kinds"): + DatasetSchema.model_validate_json( + json.dumps( + { + "observation": {"state": snapshot.model_dump()}, + "action": {"action": updates.model_dump()}, + } + ) + ) + + +def test_snapshots_can_align_forward_but_targets_remain_causal(): + store = _FakeStore( + { + "anchor": _scalar_stream([(10, 0), (11, 0)]), + "measured": _scalar_stream([(10.01, 3), (10.99, 4)]), + "commands": [ + _Obs(1, JointState(name=["left", "right", "gripper"], position=[1, 2, 0])), + _Obs(10.001, JointState(name=["left"], position=[5])), + ], + } + ) + streams = { + "anchor": _feature("anchor"), + "measured": _feature("measured"), + "target": _held_action(), + } + episode = Episode(id="episode", start_ts=10, end_ts=11) + sync = SyncConfig(anchor="anchor", rate_hz=1, tolerance_ms=20) + + report = inspect_episode_quality(store, episode, streams, sync, QualityConfig()) + samples = list( + iter_episode_samples( + store, episode, streams, sync, QualityConfig(), action_keys={"measured"} + ) + ) + + assert report.valid, report.rejection_reasons + assert report.emitted_frames == len(samples) == 2 + assert report.max_alignment_error_ms == pytest.approx(10) + np.testing.assert_array_equal([sample.action["measured"] for sample in samples], [[3], [4]]) + np.testing.assert_array_equal( + [sample.observation["target"] for sample in samples], [[1, 2, 0], [5, 2, 0]] + ) + + +def test_shared_update_source_is_read_once_for_multiple_projections(mocker): + store = _FakeStore( + { + "anchor": _scalar_stream([(10, 0)]), + "commands": [ + _Obs(1, JointState(name=["left", "right", "gripper"], position=[1, 2, 0])) + ], + } + ) + read = mocker.spy(store, "stream") + streams = { + "anchor": _feature("anchor"), + "all": _held_action(), + "gripper": FeatureSpec( + **(_held_action().model_dump() | {"names": ["gripper"], "shape": (1,)}) + ), + } + samples = list( + iter_episode_samples( + store, + Episode(id="episode", start_ts=10, end_ts=10), + streams, + SyncConfig(anchor="anchor", rate_hz=1, tolerance_ms=20), + QualityConfig(), + ) + ) + assert read.call_args_list == [mocker.call("anchor"), mocker.call("commands")] + np.testing.assert_array_equal(samples[0].observation["all"], [1, 2, 0]) + np.testing.assert_array_equal(samples[0].observation["gripper"], [0]) + + +def test_held_action_uses_causal_recording_history_across_episodes(): + store = _FakeStore( + { + "anchor": _scalar_stream([(10.0, 0.0), (11.0, 0.0), (12.0, 0.0), (20.0, 0.0)]), + "commands": [ + _Obs(1.0, JointState(name=["right", "left"], position=[2.0, 1.0])), + _Obs(2.0, JointState(name=["gripper"], position=[0.0])), + _Obs(10.0, JointState(name=["gripper"], position=[1.0])), + _Obs(10.01, JointState(name=["left"], position=[3.0])), + _Obs(12.0, JointState(name=["gripper", "right", "left"], position=[0.0, 5.0, 4.0])), + ], + } + ) + streams = {"anchor": _feature("anchor"), "action": _held_action()} + sync = SyncConfig(anchor="anchor", rate_hz=1.0, tolerance_ms=20) + for start, end, expected in [ + (10.0, 12.0, [[1, 2, 1], [3, 2, 1], [4, 5, 0]]), + (20.0, 20.0, [[4, 5, 0]]), + ]: + episode = Episode(id="episode", start_ts=start, end_ts=end) + report = inspect_episode_quality(store, episode, streams, sync, QualityConfig()) + samples = list( + iter_episode_samples( + store, + episode, + streams, + sync, + QualityConfig(), + action_keys={"action"}, + ) + ) + assert report.valid, report.rejection_reasons + assert report.emitted_frames == len(samples) == len(expected) + assert report.filled_frames == 0 + assert report.max_alignment_error_ms == 0 + np.testing.assert_array_equal([sample.action["action"] for sample in samples], expected) + + +def test_held_action_does_not_initialize_from_future_commands(): + store = _FakeStore( + { + "anchor": _scalar_stream([(10.0, 0.0), (11.0, 0.0)]), + "commands": [ + _Obs(1.0, JointState(name=["left", "right"], position=[1.0, 2.0])), + _Obs(10.001, JointState(name=["gripper"], position=[1.0])), + ], + } + ) + streams = {"anchor": _feature("anchor"), "action": _held_action()} + sync = SyncConfig(anchor="anchor", rate_hz=1, tolerance_ms=20) + episode = Episode(id="episode", start_ts=10, end_ts=11) + report = inspect_episode_quality(store, episode, streams, sync, QualityConfig()) + samples = list(iter_episode_samples(store, episode, streams, sync, QualityConfig())) + assert not report.valid + assert report.emitted_frames == len(samples) == 1 + assert "gripper" in " ".join(report.rejection_reasons) + assert "10.0" in " ".join(report.rejection_reasons) + assert samples[0].ts == 11 + + +@pytest.mark.parametrize( + "message, reason", + [ + (JointState(name=["left", "left"], position=[1, 2]), "duplicate"), + (JointState(name=["left"], position=[]), "names but"), + (JointState(name=["left"], position=[float("nan")]), "non-finite"), + ({"position": [1]}, "JointState"), + ], +) +def test_held_action_rejects_malformed_updates_even_between_samples(message, reason): + store = _FakeStore( + { + "anchor": _scalar_stream([(10.0, 0.0)]), + "commands": [ + _Obs(1, message), + _Obs(2, JointState(name=["left", "right", "gripper"], position=[1, 2, 0])), + ], + } + ) + streams = {"anchor": _feature("anchor"), "action": _held_action()} + episode = Episode(id="episode", start_ts=10, end_ts=10) + sync = SyncConfig(anchor="anchor", rate_hz=1, tolerance_ms=20) + report = inspect_episode_quality(store, episode, streams, sync, QualityConfig()) + assert not report.valid + assert reason in " ".join(report.rejection_reasons) + with pytest.raises(ValueError, match=reason): + list(iter_episode_samples(store, episode, streams, sync, QualityConfig())) + + +def test_nearest_validates_selected_samples_not_unused_partial_messages(): + store = _FakeStore( + { + "anchor": _scalar_stream([(10.0, 0.0), (11.0, 0.0)]), + "commands": [ + _Obs(10, JointState(name=["joint"], position=[1])), + _Obs(10.5, JointState(name=[], position=[])), + _Obs(11, JointState(name=["joint"], position=[2])), + ], + } + ) + streams = {"anchor": _feature("anchor"), "action": _feature("commands")} + episode = Episode(id="episode", start_ts=10, end_ts=11) + sync = SyncConfig(anchor="anchor", rate_hz=1, tolerance_ms=20) + report = inspect_episode_quality(store, episode, streams, sync, QualityConfig()) + samples = list(iter_episode_samples(store, episode, streams, sync, QualityConfig())) + assert report.valid, report.rejection_reasons + assert report.emitted_frames == len(samples) == 2 + + +def test_held_action_without_history_reports_missing_joints_and_emits_nothing(): + store = _FakeStore({"anchor": _scalar_stream([(10.0, 0.0)])}) + streams = {"anchor": _feature("anchor"), "action": _held_action()} + episode = Episode(id="episode", start_ts=10, end_ts=10) + sync = SyncConfig(anchor="anchor", rate_hz=1, tolerance_ms=20) + + report = inspect_episode_quality(store, episode, streams, sync, QualityConfig()) + + assert not report.valid + assert report.expected_frames == 1 + assert report.emitted_frames == 0 + assert "['left', 'right', 'gripper']" in report.rejection_reasons[0] + assert list(iter_episode_samples(store, episode, streams, sync, QualityConfig())) == [] + + +def test_held_actions_do_not_relax_camera_rate_or_gap_checks(): + image = {"data": np.zeros((8, 8, 3), dtype=np.uint8)} + store = _FakeStore( + { + "camera": [_Obs(10, image), _Obs(10.1, image), _Obs(10.3, image)], + "commands": [ + _Obs( + 1, + JointState( + name=["left", "right", "gripper"], + position=[1, 2, 0], + ), + ) + ], + } + ) + streams = { + "camera": FeatureSpec( + stream="camera", + field="data", + dtype="video", + shape=(8, 8, 3), + names=["height", "width", "channels"], + ), + "action": _held_action(), + } + episode = Episode(id="episode", start_ts=10, end_ts=10.3) + sync = SyncConfig(anchor="camera", rate_hz=30, tolerance_ms=20) + + report = inspect_episode_quality(store, episode, streams, sync, QualityConfig()) + samples = list(iter_episode_samples(store, episode, streams, sync, QualityConfig())) + + assert not report.valid + assert report.source_rates_hz["camera"] == pytest.approx(2 / 0.3) + assert report.max_gaps_ms["camera"] == pytest.approx(200) + reasons = " ".join(report.rejection_reasons) + assert "source rate" in reasons + assert "maximum source gap" in reasons + assert report.emitted_frames == len(samples) + + +@pytest.mark.parametrize( + "changes", + [ + {"field": "velocity"}, + {"field": None}, + {"dtype": "video", "shape": (2, 2, 3)}, + {"names": ["left", "left", "gripper"]}, + ], +) +def test_held_action_requires_named_position_vector(changes): + values = dict( + stream="commands", + field="position", + dtype="float32", + shape=(3,), + names=["left", "right", "gripper"], + source_kind="joint_position_updates", + ) + with pytest.raises(ValueError): + FeatureSpec(**(values | changes)) # ── summarize_lengths ──────────────────────────────────────────────────────── @@ -350,18 +795,18 @@ def test_summarize_lengths_empty() -> None: # ── dimos_meta sidecar ─────────────────────────────────────────────────────── -def test_dimos_meta_records_sync_and_action_shift(tmp_path: Path) -> None: +def test_dimos_meta_records_sync_and_quality(tmp_path: Path) -> None: cfg = DataPrepConfig( source="s.db", - observation={"state": StreamField(stream="js", field="position")}, - action={"action": StreamField(stream="js", field="position")}, - sync=SyncConfig(anchor="state", rate_hz=14.0, tolerance_ms=80.0, action_shift=0), + observation={"state": _feature("js")}, + action={"action": _feature("js")}, + sync=SyncConfig(anchor="state", rate_hz=14.0, tolerance_ms=80.0), output=OutputConfig(format="lerobot", path=tmp_path, metadata={"fps": 14}), ) - _write_dimos_meta(tmp_path, cfg, episodes=[]) + _write_dimos_meta(tmp_path, cfg, episodes=[], quality_reports=[]) meta = json.loads((tmp_path / "dimos_meta.json").read_text()) - assert meta["sync"]["action_shift"] == 0 + assert meta["quality"]["mode"] == "strict" assert meta["source"] == "s.db" @@ -372,7 +817,7 @@ def test_dimos_meta_beside_file_for_hdf5(tmp_path: Path) -> None: ds_file.write_bytes(b"\x89HDF\r\n") # stand-in for a real .hdf5 cfg = DataPrepConfig(source="s.db", output=OutputConfig(format="hdf5", path=ds_file)) - _write_dimos_meta(ds_file, cfg, episodes=[]) + _write_dimos_meta(ds_file, cfg, episodes=[], quality_reports=[]) sidecar = tmp_path / "session.dimos_meta.json" assert sidecar.exists() # beside the file, not session.hdf5/dimos_meta.json @@ -384,8 +829,100 @@ def test_run_dataprep_rejects_shared_obs_action_key() -> None: two maps merge; run_dataprep must reject it before opening the store.""" cfg = DataPrepConfig( source="nonexistent.db", # never reached — the check runs first - observation={"joints": StreamField(stream="joint_state", field="position")}, - action={"joints": StreamField(stream="joint_state", field="position")}, + observation={"joints": _feature("joint_state")}, + action={"joints": _feature("joint_state")}, ) with pytest.raises(ValueError, match="share feature name"): run_dataprep(cfg) + + +def test_run_dataprep_rejects_empty_recorded_stream_before_writer(mocker, tmp_path: Path) -> None: + store = mocker.MagicMock() + store.list_streams.return_value = ["color_image", "joint_state", "status"] + stream_counts = {"color_image": 0, "joint_state": 20} + store.stream.side_effect = lambda name: mocker.MagicMock( + count=mocker.Mock(return_value=stream_counts[name]) + ) + mocker.patch("dimos.imitation.dataprep.build._open_recording", return_value=store) + mocker.patch( + "dimos.imitation.dataprep.build.extract_episodes", + return_value=[Episode(id="ep_0", start_ts=1.0, end_ts=2.0)], + ) + writer = mocker.Mock(return_value=tmp_path) + cfg = DataPrepConfig( + source="recording.db", + observation={ + "wrist": FeatureSpec( + stream="color_image", + field="data", + dtype="video", + shape=(2, 2, 3), + names=["height", "width", "channels"], + ), + "state": _feature("joint_state"), + }, + output=OutputConfig(format="hdf5", path=tmp_path / "dataset.hdf5"), + sync=SyncConfig(anchor="wrist", rate_hz=30.0, tolerance_ms=20.0), + ) + + with pytest.raises(RuntimeError, match="stream 'color_image' has no episode data"): + run_dataprep(cfg, writer=writer) + + writer.assert_not_called() + store.stop.assert_called_once_with() + + +def test_run_dataprep_excludes_only_invalid_episode(mocker, tmp_path: Path) -> None: + store = mocker.MagicMock() + store.list_streams.return_value = ["joint_state", "status"] + mocker.patch("dimos.imitation.dataprep.build._open_recording", return_value=store) + episodes = [ + Episode(id="bad", start_ts=0.0, end_ts=1.0, task_label="pick"), + Episode(id="good", start_ts=2.0, end_ts=3.0, task_label="pick"), + ] + mocker.patch("dimos.imitation.dataprep.build.extract_episodes", return_value=episodes) + mocker.patch( + "dimos.imitation.dataprep.build.inspect_episode_quality", + side_effect=[ + EpisodeQualityReport( + episode_id="bad", + valid=False, + mode="strict", + rejection_reasons=["missing frame"], + ), + EpisodeQualityReport(episode_id="good", valid=True, mode="strict"), + ], + ) + sample = Sample( + ts=2.0, + episode_id="good", + observation={"state": np.asarray([1.0], dtype=np.float32)}, + action={}, + task_label="pick", + complementary_info={"is_filled": np.asarray([False])}, + ) + aligned = mocker.patch( + "dimos.imitation.dataprep.build.iter_episode_samples", return_value=iter([sample]) + ) + received: list[Sample] = [] + + def writer(samples, _output): + received.extend(samples) + return tmp_path + + config = DataPrepConfig( + source="recording.db", + observation={"state": _feature("joint_state")}, + sync=SyncConfig(anchor="state", rate_hz=1.0, tolerance_ms=20.0), + output=OutputConfig(format="hdf5", path=tmp_path / "dataset.hdf5"), + ) + + run_dataprep(config, writer=writer) + + assert [value.episode_id for value in received] == ["good"] + assert aligned.call_args.kwargs["episode"].id == "good" + reports = json.loads((tmp_path / "dimos_meta.json").read_text())["quality_reports"] + assert [(report["episode_id"], report["valid"]) for report in reports] == [ + ("bad", False), + ("good", True), + ] diff --git a/dimos/imitation/dataprep/test_lerobot.py b/dimos/imitation/dataprep/test_lerobot.py new file mode 100644 index 0000000000..8e985b222d --- /dev/null +++ b/dimos/imitation/dataprep/test_lerobot.py @@ -0,0 +1,123 @@ +# 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. + +import json +from pathlib import Path +import subprocess + +import pytest +import pytest_mock + +from dimos.imitation.dataprep.core import DataPrepConfig, OutputConfig +from dimos.imitation.dataprep.lerobot import ( + inspect_lerobot_dataset, + lerobot_project, + run_lerobot_dataprep, +) +from dimos.utils.data import get_project_root + + +def test_conversion_runs_module_in_checkout_policy_project( + tmp_path: Path, mocker: pytest_mock.MockerFixture +) -> None: + mocker.patch.dict( + "os.environ", + { + "VIRTUAL_ENV": "/parent/.venv", + "UV_PYTHON": "3.10", + "UV_PROJECT_ENVIRONMENT": "/parent/env", + }, + ) + run = mocker.patch( + "dimos.imitation.dataprep.lerobot.subprocess.run", + return_value=subprocess.CompletedProcess( + [], + 0, + stdout=( + f'[dataprep] wrote 1 episode\n{{"command":"build","path":"{tmp_path / "dataset"}"}}' + ), + stderr="", + ), + ) + config = DataPrepConfig( + source="recording.db", + output=OutputConfig(format="lerobot", path=tmp_path / "dataset"), + ) + + assert run_lerobot_dataprep(config) == tmp_path / "dataset" + + command = run.call_args.args[0] + project = get_project_root() / "native/python/lerobot" + assert lerobot_project() == project + assert command[:3] == ["uv", "run", "--frozen"] + assert command[-3:] == ["python", "-m", "dimos_lerobot.dataprep"] + assert "--python" not in command + assert run.call_args.kwargs["cwd"] == project + assert "VIRTUAL_ENV" not in run.call_args.kwargs["env"] + assert "UV_PYTHON" not in run.call_args.kwargs["env"] + assert run.call_args.kwargs["env"]["UV_PROJECT_ENVIRONMENT"] != "/parent/env" + assert run.call_args.kwargs["capture_output"] is True + assert run.call_args.kwargs["text"] is True + assert '"command":"build"' in run.call_args.kwargs["input"] + assert json.loads(run.call_args.kwargs["input"])["config"]["source"] == str( + Path("recording.db").resolve() + ) + + +def test_conversion_reports_missing_uv(tmp_path: Path, mocker: pytest_mock.MockerFixture) -> None: + mocker.patch( + "dimos.imitation.dataprep.lerobot.subprocess.run", + side_effect=FileNotFoundError("uv"), + ) + config = DataPrepConfig( + source="recording.db", + output=OutputConfig(format="lerobot", path=tmp_path / "dataset"), + ) + + with pytest.raises(RuntimeError, match="uv is required"): + run_lerobot_dataprep(config) + + +def test_conversion_reports_child_process_diagnostics( + tmp_path: Path, mocker: pytest_mock.MockerFixture +) -> None: + mocker.patch( + "dimos.imitation.dataprep.lerobot.subprocess.run", + return_value=subprocess.CompletedProcess( + [], 9, stdout="partial output", stderr="bad config" + ), + ) + config = DataPrepConfig( + source="recording.db", + output=OutputConfig(format="lerobot", path=tmp_path / "dataset"), + ) + + with pytest.raises(RuntimeError, match="status 9: bad config"): + run_lerobot_dataprep(config) + + +def test_inspection_uses_the_same_isolated_entrypoint( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, mocker: pytest_mock.MockerFixture +) -> None: + monkeypatch.chdir(tmp_path) + run = mocker.patch( + "dimos.imitation.dataprep.lerobot.subprocess.run", + return_value=subprocess.CompletedProcess( + [], 0, stdout='{"command":"inspect","info":{"format":"lerobot"}}', stderr="" + ), + ) + + assert inspect_lerobot_dataset(Path("dataset")) == {"format": "lerobot"} + assert '"command":"inspect"' in run.call_args.kwargs["input"] + assert json.loads(run.call_args.kwargs["input"])["path"] == str(tmp_path / "dataset") diff --git a/dimos/imitation/dataprep/test_mcap_source.py b/dimos/imitation/dataprep/test_mcap_source.py new file mode 100644 index 0000000000..112ba66310 --- /dev/null +++ b/dimos/imitation/dataprep/test_mcap_source.py @@ -0,0 +1,255 @@ +# 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. + +"""DataPrep interoperability with self-describing native MCAP recordings.""" + +from __future__ import annotations + +from collections.abc import Iterator +from pathlib import Path +from typing import Any + +from mcap.writer import Writer as McapWriter +import numpy as np +import pytest + +from dimos.imitation.dataprep.build import _open_recording, inspect_recording, run_dataprep +from dimos.imitation.dataprep.core import ( + DataPrepConfig, + FeatureSpec, + OutputConfig, + Sample, + SyncConfig, +) +from dimos.memory.codecs.jpeg import JpegCodec +from dimos.memory.codecs.lcm import LcmCodec +from dimos.memory.codecs.lz4 import Lz4Codec +from dimos.msgs.imitation_msgs.EpisodeStatus import EpisodeStatus +from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.msgs.sensor_msgs.JointState import JointState + + +def _register_channel( + writer: McapWriter, + name: str, + payload_type: type[Any], + message_encoding: str = "lcm", +) -> int: + return writer.register_channel( + topic=name, + message_encoding=message_encoding, + schema_id=0, + metadata={ + "dimos.payload_type": f"{payload_type.__module__}.{payload_type.__qualname__}", + "dimos.observation_time": "publish_time", + }, + ) + + +def _write_message(writer: McapWriter, channel_id: int, ts: float, message: Any) -> None: + timestamp_ns = round(ts * 1_000_000_000) + writer.add_message( + channel_id=channel_id, + log_time=timestamp_ns, + publish_time=timestamp_ns, + data=message if isinstance(message, bytes) else message.lcm_encode(), + ) + + +def _write_collection(path: Path) -> None: + with path.open("wb") as output: + writer = McapWriter(output) + writer.start(profile="dimos", library="test") + channels = { + "color_image": _register_channel(writer, "color_image", Image, "jpeg"), + "coordinator_joint_state": _register_channel( + writer, "coordinator_joint_state", JointState + ), + "applied_joint_position_command": _register_channel( + writer, "applied_joint_position_command", JointState + ), + "status": _register_channel(writer, "status", EpisodeStatus), + } + _write_message( + writer, + channels["status"], + 10.0, + EpisodeStatus( + ts=10.0, + state="recording", + last_event="start", + episodes_saved=0, + episodes_discarded=0, + task_label="pick", + ), + ) + for index in range(3): + ts = 10.0 + index / 30.0 + _write_message( + writer, + channels["color_image"], + ts, + JpegCodec().encode( + Image( + ts=ts, + frame_id="wrist_camera_link", + format=ImageFormat.RGB, + data=np.full((8, 8, 3), index, dtype=np.uint8), + ) + ), + ) + _write_message( + writer, + channels["coordinator_joint_state"], + ts, + JointState( + ts=ts, + frame_id="coordinator", + name=["shoulder", "wrist"], + position=[float(index), float(index + 1)], + velocity=[0.0, 0.0], + effort=[0.0, 0.0], + ), + ) + _write_message( + writer, + channels["applied_joint_position_command"], + ts, + JointState( + ts=ts, + frame_id="coordinator", + name=["shoulder", "wrist"], + position=[float(index) + 0.25, float(index) + 1.25], + velocity=[], + effort=[], + ), + ) + _write_message( + writer, + channels["status"], + 10.0 + 2 / 30.0, + EpisodeStatus( + ts=10.0 + 2 / 30.0, + state="idle", + last_event="save", + episodes_saved=1, + episodes_discarded=0, + task_label="pick", + ), + ) + writer.finish() + + +def _config(source: Path, output: Path) -> DataPrepConfig: + names = ["shoulder", "wrist"] + return DataPrepConfig( + source=str(source), + observation={ + "observation.images.wrist": FeatureSpec( + stream="color_image", + field="data", + dtype="video", + shape=(8, 8, 3), + names=["height", "width", "channels"], + ), + "observation.state": FeatureSpec( + stream="coordinator_joint_state", + field="position", + dtype="float32", + shape=(2,), + names=names, + ), + }, + action={ + "action": FeatureSpec( + stream="applied_joint_position_command", + field="position", + dtype="float32", + shape=(2,), + names=names, + ) + }, + sync=SyncConfig( + anchor="observation.images.wrist", + rate_hz=30.0, + tolerance_ms=20.0, + ), + output=OutputConfig(format="hdf5", path=output), + ) + + +def test_mcap_recording_inspects_and_produces_valid_samples(tmp_path: Path) -> None: + source = tmp_path / "session.mcap" + output = tmp_path / "dataset" + output.mkdir() + _write_collection(source) + config = _config(source, output) + received: list[Sample] = [] + + def writer(samples: Iterator[Sample], selected_output: OutputConfig) -> Path: + received.extend(samples) + return selected_output.path + + info = inspect_recording(source, config=config) + dataset_path = run_dataprep(config, writer=writer) + + assert info["streams"] == { + "applied_joint_position_command": 3, + "color_image": 3, + "coordinator_joint_state": 3, + "status": 2, + } + assert info["saved_episodes"] == 1 + assert info["quality"][0]["valid"] is True + assert dataset_path == output + assert len(received) == 3 + np.testing.assert_array_equal(received[0].observation["observation.state"], [0.0, 1.0]) + np.testing.assert_array_equal(received[0].action["action"], [0.25, 1.25]) + + +@pytest.mark.parametrize("encoding", ["lcm", "lz4+lcm"]) +def test_recording_metadata_decodes_typed_messages(tmp_path: Path, encoding: str) -> None: + path = tmp_path / "recording.mcap" + message = JointState(ts=12.5, name=["arm/joint1"], position=[0.25]) + codec = LcmCodec(JointState) + payload = Lz4Codec(codec).encode(message) if encoding == "lz4+lcm" else codec.encode(message) + with path.open("wb") as file: + writer = McapWriter(file) + writer.start() + channel = _register_channel(writer, "measured", JointState, encoding) + _write_message(writer, channel, 11.5, payload) + writer.finish() + + with _open_recording(path) as store: + observation = store.stream("measured").first() + assert observation.ts == 11.5 + assert observation.data.lcm_encode() == message.lcm_encode() + + +def test_missing_message_package_reports_the_recorded_type(tmp_path: Path) -> None: + path = tmp_path / "recording.mcap" + with path.open("wb") as file: + writer = McapWriter(file) + writer.start() + channel = writer.register_channel( + topic="custom_state", + message_encoding="lcm", + schema_id=0, + metadata={"dimos.payload_type": "missing_recording_package.State"}, + ) + _write_message(writer, channel, 1.0, b"unused") + writer.finish() + + with pytest.raises(ImportError, match="custom_state.*missing_recording_package.State"): + _open_recording(path) diff --git a/dimos/imitation/dataprep/test_recording.py b/dimos/imitation/dataprep/test_recording.py new file mode 100644 index 0000000000..d8254da3d8 --- /dev/null +++ b/dimos/imitation/dataprep/test_recording.py @@ -0,0 +1,255 @@ +# 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. + + +from contextlib import contextmanager + +import h5py +from mcap.writer import Writer as McapWriter +import numpy as np +import pytest + +from dimos.imitation.collection.profile import CollectionFeature, CollectionProfile +from dimos.imitation.collection.recording import RecordingSchema +from dimos.imitation.dataprep.build import inspect_recording, run_dataprep +from dimos.imitation.dataprep.core import FeatureSpec, OutputConfig, SyncConfig +from dimos.memory.store.sqlite import SqliteStore +from dimos.msgs.imitation_msgs.EpisodeStatus import EpisodeStatus +from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.msgs.sensor_msgs.JointState import JointState + + +@contextmanager +def _raw_writer(path, types): + if path.suffix == ".db": + with SqliteStore(path=str(path)) as store: + streams = {name: store.stream(name, kind, codec="lcm") for name, kind in types.items()} + yield lambda name, message: streams[name].append(message, ts=message.ts) + else: + with path.open("wb") as file: + writer = McapWriter(file) + writer.start() + channels = { + name: writer.register_channel( + topic=name, + message_encoding="lcm", + schema_id=0, + metadata={ + "dimos.payload_type": f"{kind.__module__}.{kind.__qualname__}", + "dimos.observation_time": "publish_time", + }, + ) + for name, kind in types.items() + } + + def append(name, message): + stamp = round(message.ts * 1_000_000_000) + writer.add_message( + channel_id=channels[name], + log_time=stamp, + publish_time=stamp, + data=message.lcm_encode(), + ) + + try: + yield append + finally: + writer.finish() + + +@pytest.mark.parametrize("format", ["mcap", "db"]) +@pytest.mark.parametrize("camera_count", [1, 2, 4]) +def test_moved_recording_directory_prepares_saved_episodes(format, camera_count, tmp_path): + cameras = { + f"camera_{index}": CollectionFeature( + stream=f"view_{index}", + message_type=Image, + field="data", + dtype="video", + shape=(8, 8, 3), + names=["height", "width", "channels"], + ) + for index in range(camera_count) + } + joints = CollectionFeature( + stream="measured", + message_type=JointState, + field="position", + dtype="float32", + shape=(2,), + names=["right", "left"], + ) + profile = CollectionProfile( + name="external-robot", + robot_type="custom", + observations={**cameras, "state": joints}, + actions={"action": joints}, + sync=SyncConfig(anchor="camera_0", rate_hz=30, tolerance_ms=20), + ) + directory = tmp_path / "original" + directory.mkdir() + schema = profile.to_schema() + schema.payload = f"recording.{format}" + (directory / "schema.json").write_text(schema.model_dump_json()) + types = {**profile.input_types(), "status": EpisodeStatus} + with _raw_writer(directory / schema.payload, types) as append: + for start, event in [(10.0, "save"), (20.0, "discard"), (30.0, None)]: + append( + "status", + EpisodeStatus( + task_label="pick", + ts=start, + state="recording", + last_event="start", + episodes_saved=0, + episodes_discarded=0, + ), + ) + for index in range(3): + ts = start + index / 30 + append( + "measured", + JointState(ts=ts, name=["left", "right"], position=[index, index + 10]), + ) + for camera in cameras.values(): + append( + camera.stream, + Image( + ts=ts, format=ImageFormat.RGB, data=np.zeros((8, 8, 3), dtype=np.uint8) + ), + ) + if event is not None: + append( + "status", + EpisodeStatus( + task_label="pick", + ts=ts, + state="idle", + last_event=event, + episodes_saved=1, + episodes_discarded=0, + ), + ) + moved = tmp_path / "moved" + directory.rename(moved) + restored = RecordingSchema.read(moved) + config = restored.dataprep_config( + moved, OutputConfig(format="hdf5", path=tmp_path / "dataset.h5") + ) + result = run_dataprep(config) + with h5py.File(result, "r") as dataset: + assert list(dataset["episodes"]) == ["episode_000000"] + np.testing.assert_array_equal( + dataset["episodes"]["episode_000000"]["observation"]["state"][:], + [[10, 0], [11, 1], [12, 2]], + ) + assert config.source == str(moved / f"recording.{format}") + assert restored.observation == schema.observation + + +@pytest.mark.parametrize("format", ["mcap", "db"]) +def test_sparse_actions_inspect_and_convert_with_recording_history(format, tmp_path): + schema = RecordingSchema( + name="held-positions", + robot_type="custom", + payload=f"recording.{format}", + observation={ + "state": FeatureSpec( + stream="measured", + field="position", + dtype="float32", + shape=(3,), + names=["left", "right", "gripper"], + ) + }, + action={ + "action": FeatureSpec( + stream="commands", + field="position", + dtype="float32", + shape=(3,), + names=["left", "right", "gripper"], + source_kind="joint_position_updates", + ) + }, + sync=SyncConfig(anchor="state", rate_hz=1, tolerance_ms=20), + ) + (tmp_path / "schema.json").write_text(schema.model_dump_json()) + source = tmp_path / schema.payload + with _raw_writer( + source, + { + "measured": JointState, + "commands": JointState, + "status": EpisodeStatus, + }, + ) as append: + # Deliberately insert updates out of timestamp order. + append("commands", JointState(ts=2, name=["gripper"], position=[0])) + append("commands", JointState(ts=1, name=["right", "left"], position=[2, 1])) + append("commands", JointState(ts=11, name=["gripper"], position=[1])) + for index, start in enumerate([10, 20]): + append( + "status", + EpisodeStatus( + ts=start, + state="recording", + last_event="start", + episodes_saved=index, + episodes_discarded=0, + task_label="fold", + ), + ) + for ts in [start, start + 1]: + append( + "measured", + JointState( + ts=ts, + name=["left", "right", "gripper"], + position=[0, 0, 0], + ), + ) + append( + "status", + EpisodeStatus( + ts=start + 1, + state="idle", + last_event="save", + episodes_saved=index + 1, + episodes_discarded=0, + task_label="fold", + ), + ) + restored = RecordingSchema.read(tmp_path) + config = restored.dataprep_config( + tmp_path, + OutputConfig(format="hdf5", path=tmp_path / "dataset.h5"), + ) + + info = inspect_recording(source, config=config) + result = run_dataprep(config) + + assert [ + (report["valid"], report["emitted_frames"], report["filled_frames"]) + for report in info["quality"] + ] == [(True, 2, 0), (True, 2, 0)] + with h5py.File(result, "r") as dataset: + np.testing.assert_array_equal( + dataset["episodes/episode_000000/action/action"], + [[1, 2, 0], [1, 2, 1]], + ) + np.testing.assert_array_equal( + dataset["episodes/episode_000001/action/action"], + [[1, 2, 1], [1, 2, 1]], + ) diff --git a/dimos/imitation/test_datacollection_e2e.py b/dimos/imitation/test_datacollection_e2e.py index 7c931ce403..39cd243961 100644 --- a/dimos/imitation/test_datacollection_e2e.py +++ b/dimos/imitation/test_datacollection_e2e.py @@ -12,78 +12,56 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""End-to-end coverage from live collection through both dataset formats.""" +"""End-to-end coverage from live collection through host-side DataPrep.""" from __future__ import annotations -from collections.abc import Callable import json from pathlib import Path -from typing import Any +import subprocess +from typing import Any, cast +import uuid -import cv2 import h5py import numpy as np -import pyarrow.parquet as pq import pytest -from dimos.core.stream import Stream, Transport -from dimos.imitation.collection.episode_monitor import ( - EpisodeEvent, - EpisodeStatus, - RecordingState, -) -from dimos.imitation.collection.recorder import CollectionRecorder +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.profile import CollectionFeature, CollectionProfile +from dimos.imitation.collection.recorder import collection_recorder +from dimos.imitation.collection.recording import RecordingSchema from dimos.imitation.dataprep.build import inspect_dataset, run_dataprep from dimos.imitation.dataprep.core import ( DataPrepConfig, EpisodeExtractor, + FeatureSpec, OutputConfig, - StreamField, + QualityConfig, SyncConfig, extract_episodes, ) from dimos.memory.store.sqlite import SqliteStore +from dimos.msgs.imitation_msgs.EpisodeStatus import ( + EpisodeEvent, + EpisodeStatus, + RecordingState, +) +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.protocol.pubsub.impl.zenohpubsub import QOS_NEVER_DROP, Topic 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, ] -class _DirectTransport(Transport[Any]): - """Synchronous in-process transport used to exercise real port subscriptions.""" - - def __init__(self) -> None: - self._subscribers: list[Callable[[Any], Any]] = [] - - def start(self) -> None: - pass - - def stop(self) -> None: - self._subscribers.clear() - - def broadcast(self, selfstream: Stream[Any] | None, value: Any) -> None: - for callback in tuple(self._subscribers): - callback(value) - - def subscribe( - self, - callback: Callable[[Any], Any], - selfstream: Stream[Any] | None = None, - ) -> Callable[[], None]: - self._subscribers.append(callback) - - def unsubscribe() -> None: - self._subscribers.remove(callback) - - return unsubscribe - - def _status( ts: float, event: EpisodeEvent, @@ -107,37 +85,87 @@ def _dataprep_config(db_path: Path, output: OutputConfig) -> DataPrepConfig: source=str(db_path), episodes=EpisodeExtractor(status_stream="status"), observation={ - "camera": StreamField(stream="color_image"), - "state": StreamField(stream="coordinator_joint_state", field="position"), + "camera": FeatureSpec( + stream="color_image", + field="data", + dtype="video", + shape=(16, 16, 3), + names=["height", "width", "channels"], + ), + "state": FeatureSpec( + stream="coordinator_joint_state", + field="position", + dtype="float32", + shape=(2,), + names=["joint_0", "joint_1"], + ), }, action={ - "action": StreamField(stream="coordinator_joint_state", field="position"), + "action": FeatureSpec( + stream="coordinator_joint_state", + field="position", + dtype="float32", + shape=(2,), + names=["joint_0", "joint_1"], + ), }, - sync=SyncConfig(anchor="camera", rate_hz=1.0, tolerance_ms=1.0, action_shift=1), + sync=SyncConfig(anchor="camera", rate_hz=1.0, tolerance_ms=1.0), + quality=QualityConfig(max_camera_gap_ms=1100.0), output=output, ) -def _record_session(db_path: Path) -> None: - recorder = CollectionRecorder( - db_path=db_path, - record_tf=False, - poseless_streams=["color_image", "coordinator_joint_state", "status"], +def _record_session(db_path: Path, executable: Path) -> dict[str, int]: + config = _dataprep_config(db_path, OutputConfig(path=db_path.parent / "unused")) + profile = CollectionProfile( + name="synthetic", + robot_type="synthetic", + observations={ + name: CollectionFeature( + **feature.model_dump(), + message_type=Image if name == "camera" else JointState, + ) + for name, feature in config.observation.items() + }, + actions={ + name: CollectionFeature(**feature.model_dump(), message_type=JointState) + for name, feature in config.action.items() + }, + sync=config.sync, + quality=config.quality, ) + atom = collection_recorder( + profile=profile, recording=db_path.parent, format="sqlite" + ).active_blueprints[0] + recorder = atom.module(**atom.kwargs, executable=str(executable)) + topic_prefix = f"dimos/test/collection-export/{uuid.uuid4().hex}" + payload_types = { + "color_image": Image, + "coordinator_joint_state": JointState, + "status": EpisodeStatus, + } transports = { - "color_image": _DirectTransport(), - "coordinator_joint_state": _DirectTransport(), - "status": _DirectTransport(), + name: ZenohTransport( + Topic(f"{topic_prefix}/{name}", cast("type[DimosMsg]", kind), qos=QOS_NEVER_DROP) + ) + for name, kind in payload_types.items() } for name, transport in transports.items(): getattr(recorder, name).transport = transport counts = {name: 0 for name in transports} + def stream_count(name: str) -> int: + with SqliteStore(path=str(db_path), must_exist=True) as store: + stream = store.stream(name) + if name == "status": + return sum(obs.data.last_event != "init" for obs in stream.to_list()) + return stream.count() + def publish(name: str, message: Any) -> None: counts[name] += 1 - transports[name].publish(message) + transports[name].broadcast(None, message) wait_until( - lambda: recorder.store.stream(name).count() == counts[name], + lambda: stream_count(name) == counts[name], timeout=5.0, interval=0.005, message=f"{name} message {counts[name]} was not recorded", @@ -150,6 +178,21 @@ def publish(name: str, message: Any) -> None: ] try: recorder.start() + ready = _status(1.0, "init", "idle", 0, 0, "") + + def received_probe() -> bool: + with SqliteStore(path=str(db_path), must_exist=True) as store: + if store.stream("status").count() > 0: + return True + transports["status"].broadcast(None, ready) + return False + + wait_until( + received_probe, + timeout=10.0, + interval=0.1, + message="native collection status subscription did not become ready", + ) saved = 0 discarded = 0 for start_ts, task, success, base in episodes: @@ -193,41 +236,47 @@ def publish(name: str, message: Any) -> None: publish("status", _status(112.0, "start", "recording", 2, 1, "interrupted")) finally: recorder.stop() - - -def _read_video(path: Path) -> list[np.ndarray[Any, Any]]: - capture = cv2.VideoCapture(str(path)) - frames: list[np.ndarray[Any, Any]] = [] - try: - while True: - ok, bgr = capture.read() - if not ok: - return frames - frames.append(cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)) - finally: - capture.release() + for transport in transports.values(): + transport.stop() + return counts EXPECTED_STATE = np.asarray( - [[0.0, 100.0], [1.0, 101.0], [20.0, 120.0], [21.0, 121.0]], - dtype=np.float32, -) -EXPECTED_ACTION = np.asarray( - [[1.0, 101.0], [2.0, 102.0], [21.0, 121.0], [22.0, 122.0]], + [ + [0.0, 100.0], + [1.0, 101.0], + [2.0, 102.0], + [20.0, 120.0], + [21.0, 121.0], + [22.0, 122.0], + ], dtype=np.float32, ) +EXPECTED_ACTION = EXPECTED_STATE.copy() @pytest.fixture(scope="module") def recorded_session( tmp_path_factory: pytest.TempPathFactory, ) -> tuple[Path, dict[float, np.ndarray[Any, Any]]]: - db_path = tmp_path_factory.mktemp("recorded-session") / "recording.db" - _record_session(db_path) + subprocess.run( + ["cargo", "build", "--locked", "-p", "dimos-memory-recorder"], + cwd=DIMOS_PROJECT_ROOT, + check=True, + ) + executable = DIMOS_PROJECT_ROOT / "target" / "debug" / "dimos-memory-recorder" + db_path = tmp_path_factory.mktemp("recorded-session") / "session" / "recording.db" + with pytest.MonkeyPatch.context() as patch: + patch.setattr(global_config, "transport", "zenoh") + counts = _record_session(db_path, executable) with SqliteStore(path=str(db_path), must_exist=True) as store: assert store.stream("color_image").count() == 9 assert store.stream("coordinator_joint_state").count() == 9 - assert store.stream("status").count() == 7 + assert ( + sum(obs.data.last_event != "init" for obs in store.stream("status").to_list()) + == counts["status"] + == 7 + ) episodes = extract_episodes(store, EpisodeExtractor(status_stream="status")) assert [ (episode.start_ts, episode.end_ts, episode.success, episode.task_label) @@ -259,8 +308,8 @@ def test_collection_to_hdf5_roundtrip( db_path, recorded_images = recorded_session hdf5_path = run_dataprep( - _dataprep_config( - db_path, + RecordingSchema.read(db_path.parent).dataprep_config( + db_path.parent, OutputConfig( format="hdf5", path=tmp_path / "dataset.hdf5", @@ -269,11 +318,11 @@ def test_collection_to_hdf5_roundtrip( ) ) hdf5_info = inspect_dataset(hdf5_path) - assert (hdf5_info["episodes"], hdf5_info["frames"], hdf5_info["fps"]) == (2, 4, 1.0) + assert (hdf5_info["episodes"], hdf5_info["frames"], hdf5_info["fps"]) == (2, 6, 1.0) assert hdf5_info["episode_lengths"] == { - "min": 2, - "max": 2, - "mean": 2.0, + "min": 3, + "max": 3, + "mean": 3.0, "uniform": True, } @@ -281,8 +330,8 @@ def test_collection_to_hdf5_roundtrip( first = h5["episodes/episode_000000"] second = h5["episodes/episode_000001"] assert [first.attrs["start_ts"], second.attrs["start_ts"]] == [100.0, 108.0] - np.testing.assert_array_equal(first["timestamp"][:], [0.0, 1.0]) - np.testing.assert_array_equal(second["timestamp"][:], [0.0, 1.0]) + np.testing.assert_array_equal(first["timestamp"][:], [0.0, 1.0, 2.0]) + np.testing.assert_array_equal(second["timestamp"][:], [0.0, 1.0, 2.0]) np.testing.assert_array_equal( np.concatenate([first["observation/state"][:], second["observation/state"][:]]), EXPECTED_STATE, @@ -293,11 +342,11 @@ def test_collection_to_hdf5_roundtrip( ) np.testing.assert_array_equal( first["observation/camera"][:], - np.stack([recorded_images[100.0], recorded_images[101.0]]), + np.stack([recorded_images[100.0], recorded_images[101.0], recorded_images[102.0]]), ) np.testing.assert_array_equal( second["observation/camera"][:], - np.stack([recorded_images[108.0], recorded_images[109.0]]), + np.stack([recorded_images[108.0], recorded_images[109.0], recorded_images[110.0]]), ) hdf5_meta = json.loads((tmp_path / "dataset.dimos_meta.json").read_text()) @@ -305,51 +354,3 @@ def test_collection_to_hdf5_roundtrip( (episode["start_ts"], episode["end_ts"], episode["task_label"]) for episode in hdf5_meta["episodes"] ] == [(100.0, 102.0, "pick"), (108.0, 110.0, "place")] - - -def test_collection_to_lerobot_roundtrip( - tmp_path: Path, - recorded_session: tuple[Path, dict[float, np.ndarray[Any, Any]]], -) -> None: - db_path, recorded_images = recorded_session - try: - lerobot_path = run_dataprep( - _dataprep_config( - db_path, - OutputConfig( - format="lerobot", - path=tmp_path / "lerobot", - metadata={"robot": "synthetic"}, - ), - ) - ) - except RuntimeError as exc: - if "VideoWriter" in str(exc): - pytest.skip(f"no mp4v encoder available in this environment: {exc}") - raise - - lerobot_info = inspect_dataset(lerobot_path) - assert (lerobot_info["episodes"], lerobot_info["frames"], lerobot_info["fps"]) == (2, 4, 1.0) - data = pq.read_table(lerobot_path / "data/chunk-000/file-000.parquet") - assert data.column("timestamp").to_pylist() == pytest.approx([0.0, 1.0, 0.0, 1.0]) - assert data.column("episode_index").to_pylist() == [0, 0, 1, 1] - assert data.column("frame_index").to_pylist() == [0, 1, 0, 1] - np.testing.assert_array_equal( - np.asarray(data.column("observation.state").to_pylist()), EXPECTED_STATE - ) - np.testing.assert_array_equal(np.asarray(data.column("action").to_pylist()), EXPECTED_ACTION) - - episode_rows = pq.read_table( - lerobot_path / "meta/episodes/chunk-000/file-000.parquet" - ).to_pylist() - assert [row["length"] for row in episode_rows] == [2, 2] - assert [(row["dataset_from_index"], row["dataset_to_index"]) for row in episode_rows] == [ - (0, 2), - (2, 4), - ] - assert [row["tasks"] for row in episode_rows] == [["pick"], ["place"]] - - video = _read_video(lerobot_path / "videos/observation.images.camera/chunk-000/file-000.mp4") - assert len(video) == 4 - expected_means = [recorded_images[ts].mean() for ts in (100.0, 101.0, 108.0, 109.0)] - np.testing.assert_allclose([frame.mean() for frame in video], expected_means, atol=5.0) diff --git a/dimos/mapping/dim_slam/rust/Cargo.lock b/dimos/mapping/dim_slam/rust/Cargo.lock index 1610904010..8c5058b493 100644 --- a/dimos/mapping/dim_slam/rust/Cargo.lock +++ b/dimos/mapping/dim_slam/rust/Cargo.lock @@ -557,7 +557,7 @@ dependencies = [ "cu_vslam_rs", "dim_slam", "dimos-module", - "lcm-msgs", + "lcm-msgs 0.1.0 (git+https://github.com/dimensionalOS/dimos-lcm.git?branch=rust-codegen)", "serde", "tokio", "tracing", @@ -578,7 +578,7 @@ dependencies = [ [[package]] name = "dimos-lcm" version = "0.1.0" -source = "git+https://github.com/dimensionalOS/dimos-lcm.git?branch=rust-codegen#04d78e8622500244123ba9cefa4c51b4cb454549" +source = "git+https://github.com/dimensionalOS/dimos-lcm.git?rev=dd2159513ebfaa7ebc5fc32bf60209cd09aa1ca5#dd2159513ebfaa7ebc5fc32bf60209cd09aa1ca5" dependencies = [ "byteorder", "socket2 0.5.10", @@ -591,7 +591,7 @@ version = "0.1.0" dependencies = [ "dimos-lcm", "dimos-module-macros", - "lcm-msgs", + "lcm-msgs 0.1.0 (git+https://github.com/dimensionalOS/dimos-lcm.git?rev=dd2159513ebfaa7ebc5fc32bf60209cd09aa1ca5)", "nalgebra", "rayon", "serde", @@ -1393,6 +1393,14 @@ dependencies = [ "byteorder", ] +[[package]] +name = "lcm-msgs" +version = "0.1.0" +source = "git+https://github.com/dimensionalOS/dimos-lcm.git?rev=dd2159513ebfaa7ebc5fc32bf60209cd09aa1ca5#dd2159513ebfaa7ebc5fc32bf60209cd09aa1ca5" +dependencies = [ + "byteorder", +] + [[package]] name = "libc" version = "0.2.189" diff --git a/dimos/mapping/ray_tracing/rust/Cargo.toml b/dimos/mapping/ray_tracing/rust/Cargo.toml index 501f8c5028..73445f8093 100644 --- a/dimos/mapping/ray_tracing/rust/Cargo.toml +++ b/dimos/mapping/ray_tracing/rust/Cargo.toml @@ -33,7 +33,7 @@ region_bounds = "geometry_msgs.PoseStamped" [dependencies] dimos-module = { path = "../../../../native/rust/dimos-module" } -lcm-msgs = { git = "https://github.com/dimensionalOS/dimos-lcm.git", branch = "rust-codegen" } +lcm-msgs = { git = "https://github.com/dimensionalOS/dimos-lcm.git", rev = "dd2159513ebfaa7ebc5fc32bf60209cd09aa1ca5" } tokio = { version = "1", features = ["rt-multi-thread", "macros", "signal"] } serde = { version = "1", features = ["derive"] } ahash = "0.8" diff --git a/dimos/memory/module.py b/dimos/memory/module.py index 4293b561cf..50b7532cb9 100644 --- a/dimos/memory/module.py +++ b/dimos/memory/module.py @@ -445,9 +445,14 @@ def _resolve_ts(self, name: str, msg: Any) -> float: return getattr(msg, "ts", None) or time.time() async def _resolve_pose(self, name: str, msg: Any, ts: float) -> Pose | None: - """Pose to anchor *msg* with. Dispatches to the stream's (async) - ``@pose_setter_for`` if one is defined, else falls back to a - ``world <- frame_id`` tf lookup.""" + """Pose to anchor *msg* with. + + Poseless streams skip pose setters and tf resolution. Other streams + dispatch to their async ``@pose_setter_for`` when defined, then fall + back to a ``world <- frame_id`` tf lookup. + """ + if name in self.config.poseless_streams: + return None setter = self._pose_setters.get(name) if setter is not None: return cast("Pose | None", await setter(msg)) diff --git a/dimos/msgs/imitation_msgs/EpisodeStatus.py b/dimos/msgs/imitation_msgs/EpisodeStatus.py new file mode 100644 index 0000000000..ea72e95f83 --- /dev/null +++ b/dimos/msgs/imitation_msgs/EpisodeStatus.py @@ -0,0 +1,64 @@ +# 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. + +from __future__ import annotations + +from typing import ClassVar, Literal, TypeAlias, cast + +from dimos_lcm.imitation_msgs import EpisodeStatus as LCMEpisodeStatus +from pydantic import BaseModel + +from dimos.msgs.std_msgs.Header import Header + +EpisodeEvent: TypeAlias = Literal["start", "save", "discard", "init"] +RecordingState: TypeAlias = Literal["idle", "recording"] + + +class EpisodeStatus(BaseModel): + """Source-timestamped status update for an imitation-learning episode.""" + + msg_name: ClassVar[str] = "imitation_msgs.EpisodeStatus" + + ts: float + state: RecordingState + episodes_saved: int + episodes_discarded: int + last_event: EpisodeEvent = "init" + task_label: str | None = None + + def lcm_encode(self) -> bytes: + return cast( + "bytes", + LCMEpisodeStatus( + header=Header(self.ts), + state=self.state, + episodes_saved=self.episodes_saved, + episodes_discarded=self.episodes_discarded, + last_event=self.last_event, + task_label=self.task_label or "", + ).lcm_encode(), + ) + + @classmethod + def lcm_decode(cls, data: bytes) -> EpisodeStatus: + message = LCMEpisodeStatus.lcm_decode(data) + stamp = message.header.stamp + return cls( + ts=stamp.sec + stamp.nsec / 1_000_000_000, + state=message.state, + episodes_saved=message.episodes_saved, + episodes_discarded=message.episodes_discarded, + last_event=message.last_event, + task_label=message.task_label or None, + ) diff --git a/dimos/msgs/imitation_msgs/test_episode_status.py b/dimos/msgs/imitation_msgs/test_episode_status.py new file mode 100644 index 0000000000..b6c62d0e20 --- /dev/null +++ b/dimos/msgs/imitation_msgs/test_episode_status.py @@ -0,0 +1,35 @@ +# 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. + +import pytest + +from dimos.msgs.imitation_msgs.EpisodeStatus import EpisodeStatus + + +@pytest.mark.parametrize("task_label", ["pick", None]) +def test_episode_status_lcm_roundtrip_preserves_status_update( + task_label: str | None, +) -> None: + expected = EpisodeStatus( + ts=12.25, + state="recording", + episodes_saved=2, + episodes_discarded=1, + last_event="start", + task_label=task_label, + ) + + actual = EpisodeStatus.lcm_decode(expected.lcm_encode()) + + assert actual == expected diff --git a/dimos/navigation/nav_3d/mls_planner/rust/Cargo.toml b/dimos/navigation/nav_3d/mls_planner/rust/Cargo.toml index 2c5d39de04..b6e1a82fa2 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/Cargo.toml +++ b/dimos/navigation/nav_3d/mls_planner/rust/Cargo.toml @@ -36,7 +36,7 @@ path = "nav_msgs.Path" [dependencies] dimos-module = { path = "../../../../../native/rust/dimos-module" } -lcm-msgs = { git = "https://github.com/dimensionalOS/dimos-lcm.git", branch = "rust-codegen" } +lcm-msgs = { git = "https://github.com/dimensionalOS/dimos-lcm.git", rev = "dd2159513ebfaa7ebc5fc32bf60209cd09aa1ca5" } tokio = { version = "1", features = ["rt-multi-thread", "macros", "signal"] } serde = { version = "1", features = ["derive"] } ahash = "0.8" diff --git a/dimos/perception/fiducial/blueprints/desk_marker_tf.py b/dimos/perception/fiducial/blueprints/desk_marker_tf.py index 0a63f90ca5..910f09f5f7 100644 --- a/dimos/perception/fiducial/blueprints/desk_marker_tf.py +++ b/dimos/perception/fiducial/blueprints/desk_marker_tf.py @@ -26,7 +26,7 @@ from dimos.core.stream import Out from dimos.core.transport import LCMTransport from dimos.hardware.sensors.camera.module import CameraModule -from dimos.hardware.sensors.camera.webcam import Webcam +from dimos.hardware.sensors.camera.webcam import WebcamConfig from dimos.msgs.geometry_msgs.Quaternion import Quaternion from dimos.msgs.geometry_msgs.Transform import Transform from dimos.msgs.geometry_msgs.Vector3 import Vector3 @@ -43,13 +43,13 @@ DEFAULT_DESK_CAMERA_INFO_YAML = Path(__file__).resolve().parent / "fixtures" / "camera_info.yaml" -def create_desk_webcam( +def create_desk_webcam_config( camera_info_yaml: str | Path = DEFAULT_DESK_CAMERA_INFO_YAML, camera_index: int = 0, fps: float = 15.0, -) -> Webcam: +) -> WebcamConfig: camera_info = create_desk_camera_info(camera_info_yaml) - return Webcam( + return WebcamConfig( camera_index=camera_index, width=camera_info.width, height=camera_info.height, @@ -153,7 +153,7 @@ def publish_static_chain(self) -> None: desk_marker_tf = autoconnect( DeskStaticTfModule.blueprint(), CameraModule.blueprint( - hardware=create_desk_webcam, + hardware=create_desk_webcam_config(), transform=None, ), MarkerDetectionStreamModule.blueprint( diff --git a/dimos/perception/fiducial/blueprints/test_desk_marker_tf.py b/dimos/perception/fiducial/blueprints/test_desk_marker_tf.py index fc7e4ecf63..5b199f9859 100644 --- a/dimos/perception/fiducial/blueprints/test_desk_marker_tf.py +++ b/dimos/perception/fiducial/blueprints/test_desk_marker_tf.py @@ -16,7 +16,7 @@ from dimos.core.coordination.blueprints import Blueprint from dimos.hardware.sensors.camera.module import CameraModule -from dimos.hardware.sensors.camera.webcam import Webcam +from dimos.hardware.sensors.camera.webcam import WebcamConfig from dimos.perception.fiducial.blueprints.desk_marker_tf import ( DESK_CAMERA_FRAME_ID, DESK_MARKER_ARUCO_DICTIONARY, @@ -24,7 +24,7 @@ DESK_MARKER_NAMESPACE_PREFIX, DeskStaticTfModule, create_desk_camera_info, - create_desk_webcam, + create_desk_webcam_config, desk_marker_tf, ) from dimos.perception.fiducial.marker_detection_stream_module import MarkerDetectionStreamModule @@ -36,7 +36,7 @@ def test_desk_marker_tf_blueprint_declares_static_tf_module() -> None: assert isinstance(desk_marker_tf, Blueprint) assert desk_marker_tf.blueprints[0].module is DeskStaticTfModule assert desk_marker_tf.blueprints[1].module is CameraModule - assert desk_marker_tf.blueprints[1].kwargs["hardware"] is create_desk_webcam + assert isinstance(desk_marker_tf.blueprints[1].kwargs["hardware"], WebcamConfig) assert desk_marker_tf.blueprints[1].kwargs["transform"] is None assert desk_marker_tf.blueprints[2].module is MarkerDetectionStreamModule assert desk_marker_tf.blueprints[2].kwargs["marker_length_m"] == DESK_MARKER_LENGTH_M @@ -53,7 +53,7 @@ def test_desk_marker_tf_blueprint_declares_static_tf_module() -> None: ) -def test_create_desk_webcam_loads_camera_info_yaml(tmp_path: Path) -> None: +def test_create_desk_webcam_config_loads_camera_info_yaml(tmp_path: Path) -> None: camera_info_yaml = tmp_path / "camera_info.yaml" camera_info_yaml.write_text( """ @@ -80,14 +80,14 @@ def test_create_desk_webcam_loads_camera_info_yaml(tmp_path: Path) -> None: """.lstrip() ) - camera = create_desk_webcam(camera_info_yaml, camera_index=1, fps=7.5) + camera = create_desk_webcam_config(camera_info_yaml, camera_index=1, fps=7.5) - assert isinstance(camera, Webcam) - assert camera.config.camera_index == 1 - assert camera.config.width == 1920 - assert camera.config.height == 1080 - assert camera.config.fps == 7.5 - assert camera.config.camera_info.frame_id == DESK_CAMERA_FRAME_ID + assert isinstance(camera, WebcamConfig) + assert camera.camera_index == 1 + assert camera.width == 1920 + assert camera.height == 1080 + assert camera.fps == 7.5 + assert camera.camera_info.frame_id == DESK_CAMERA_FRAME_ID camera_info = create_desk_camera_info(camera_info_yaml) assert camera_info.frame_id == DESK_CAMERA_FRAME_ID diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index ed37be1220..d0bbc06441 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -101,6 +101,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", "pointlio-rust": "dimos.hardware.sensors.lidar.pointlio.pointlio_blueprints:pointlio_rust", "pointlio-rust-replay": "dimos.hardware.sensors.lidar.pointlio.pointlio_blueprints:pointlio_rust_replay", "r1pro-coordinator": "dimos.robot.galaxea.r1pro.blueprints.basic.r1pro_coordinator:r1pro_coordinator", 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 0000000000..766d7967e1 --- /dev/null +++ b/dimos/robot/manipulators/openyam/blueprints/learning_quest_collection.py @@ -0,0 +1,42 @@ +# 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 with explicit robot and camera composition.""" + +from dimos.core.coordination.blueprints import autoconnect +from dimos.hardware.sensors.camera.module import CameraModule +from dimos.hardware.sensors.camera.webcam import WebcamConfig +from dimos.imitation.collection.episode_monitor import EpisodeMonitorModule +from dimos.imitation.collection.recorder import collection_recorder +from dimos.robot.manipulators.openyam.blueprints.teleop import teleop_webxr_openyam +from dimos.robot.manipulators.openyam.collection import OPENYAM_QUEST_COLLECTION + +openyam_quest_collection = autoconnect( + teleop_webxr_openyam, + CameraModule.blueprint( + instance_name="wrist", + hardware=WebcamConfig( + camera_index=0, width=640, height=480, fps=30, frame_id_prefix="wrist_image" + ), + frame_id="wrist_camera_link", + ), + collection_recorder(profile=OPENYAM_QUEST_COLLECTION), + EpisodeMonitorModule.blueprint(instance_name="episodes"), +).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 0000000000..ca49d95f35 --- /dev/null +++ b/dimos/robot/manipulators/openyam/blueprints/test_learning_collection_e2e.py @@ -0,0 +1,220 @@ +# 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 + +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.recorder import collection_recorder +from dimos.imitation.dataprep.core import EpisodeExtractor, extract_episodes +from dimos.memory.store.sqlite import SqliteStore +from dimos.msgs.imitation_msgs.EpisodeStatus import EpisodeStatus +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.tf2_msgs.TFMessage import TFMessage +from dimos.protocol.pubsub.impl.zenohpubsub import QOS_NEVER_DROP, Topic +from dimos.robot.manipulators.openyam.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 = DIMOS_PROJECT_ROOT / "target" / "debug" / "dimos-memory-recorder" + + +@pytest.fixture(scope="module") +def native_recorder_executable() -> Path: + 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 = collection_recorder( + profile=OPENYAM_QUEST_COLLECTION, recording=directory, format="sqlite" + ).active_blueprints[0] + recorder = atom.module( + executable=str(native_recorder_executable), + **atom.kwargs, + ) + topic_prefix = f"dimos/test/native-collection/{uuid.uuid4().hex}" + payload_types: dict[str, type[Any]] = { + "wrist_image": Image, + "coordinator_joint_state": JointState, + "applied_joint_position_command": JointState, + "status": EpisodeStatus, + "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, 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( + "applied_joint_position_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, + "applied_joint_position_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("applied_joint_position_command", JointState).last().data + statuses = [ + (observation.ts, observation.data.last_event) + for observation in store.stream("status", EpisodeStatus).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/dimos/robot/manipulators/openyam/collection.py b/dimos/robot/manipulators/openyam/collection.py new file mode 100644 index 0000000000..0ff972b22e --- /dev/null +++ b/dimos/robot/manipulators/openyam/collection.py @@ -0,0 +1,66 @@ +# 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. + +"""OpenYAM collection presets, safe to import without robot hardware.""" + +from dimos.imitation.collection.profile import CollectionFeature, CollectionProfile +from dimos.imitation.dataprep.core import QualityConfig, SourceKind, SyncConfig +from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.robot.manipulators.openyam.joints import OPENYAM_JOINTS + + +def _profile( + name: str, action_stream: str, *, action_source_kind: SourceKind = "snapshot" +) -> CollectionProfile: + return CollectionProfile( + name=name, + robot_type="openyam", + observations={ + "observation.images.wrist": CollectionFeature( + stream="wrist_image", + message_type=Image, + field="data", + dtype="video", + shape=(480, 640, 3), + names=["height", "width", "channels"], + ), + "observation.state": CollectionFeature( + stream="coordinator_joint_state", + message_type=JointState, + field="position", + dtype="float32", + shape=(len(OPENYAM_JOINTS),), + names=list(OPENYAM_JOINTS), + ), + }, + actions={ + "action": CollectionFeature( + stream=action_stream, + source_kind=action_source_kind, + message_type=JointState, + field="position", + dtype="float32", + shape=(len(OPENYAM_JOINTS),), + names=list(OPENYAM_JOINTS), + ), + }, + sync=SyncConfig(anchor="observation.images.wrist", rate_hz=30, tolerance_ms=20), + quality=QualityConfig(mode="strict"), + ) + + +OPENYAM_QUEST_COLLECTION = _profile( + "openyam-quest", "applied_joint_position_command", action_source_kind="joint_position_updates" +) diff --git a/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_teleop.py b/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_teleop.py index 012b854738..04b6d28a30 100644 --- a/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_teleop.py +++ b/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_teleop.py @@ -55,11 +55,13 @@ from dimos.core.stream import In from dimos.core.transport import pSHMTransport from dimos.imitation.collection.episode_monitor import EpisodeMonitorModule -from dimos.imitation.collection.recorder import CollectionRecorder from dimos.manipulation.manipulation_module import ManipulationModule from dimos.manipulation.visualization.viser.config import ViserVisualizationConfig +from dimos.memory.module import Recorder from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.imitation_msgs.EpisodeStatus import EpisodeStatus from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.sensor_msgs.JointState import JointState from dimos.robot.unitree.g1.blueprints.basic.unitree_g1_groot_wbc import ( _G1GrootCoordinator, _unitree_g1_groot_wbc_core, @@ -68,8 +70,8 @@ from dimos.teleop.webxr.extensions import VideoArmTeleopModule -class G1CollectionRecorder(CollectionRecorder): - """CollectionRecorder plus the operator's absolute controller poses. +class G1CollectionRecorder(Recorder): + """Record G1 observations and the operator's absolute controller poses. The shared teleop IK captures controller and robot references internally, so joint commands do not appear on a stream. Recording both controller @@ -80,6 +82,9 @@ class G1CollectionRecorder(CollectionRecorder): # a GIL with control modules. dedicated_worker = True + color_image: In[Image] + coordinator_joint_state: In[JointState] + status: In[EpisodeStatus] left_cartesian_command: In[PoseStamped] right_cartesian_command: In[PoseStamped] @@ -123,6 +128,7 @@ class G1ManipulationModule(ManipulationModule): EpisodeMonitorModule.blueprint(), # default button_map: toggle=B, discard=Y G1CollectionRecorder.blueprint( db_path=_session_db(), + record_tf=False, # Collection observations/actions are synchronized by timestamp, # not localized in the world frame. Declaring them poseless also # avoids attempting a world-to-camera lookup when nav localization diff --git a/dimos/robot/unitree/g1/blueprints/primitive/unitree_g1_primitive_no_nav.py b/dimos/robot/unitree/g1/blueprints/primitive/unitree_g1_primitive_no_nav.py index 8fc4b90633..2ae1ebe65b 100644 --- a/dimos/robot/unitree/g1/blueprints/primitive/unitree_g1_primitive_no_nav.py +++ b/dimos/robot/unitree/g1/blueprints/primitive/unitree_g1_primitive_no_nav.py @@ -23,7 +23,7 @@ from dimos.core.global_config import global_config from dimos.core.transport import LCMTransport from dimos.hardware.sensors.camera.module import CameraModule -from dimos.hardware.sensors.camera.webcam import Webcam +from dimos.hardware.sensors.camera.webcam import WebcamConfig from dimos.hardware.sensors.camera.zed import compat as zed from dimos.mapping.costmapper import CostMapper from dimos.mapping.voxels.module import VoxelGridMapper @@ -97,15 +97,6 @@ def _g1_rerun_blueprint() -> Any: _with_vis = vis_module(viewer_backend=global_config.viewer, rerun_config=rerun_config) -def _create_webcam() -> Webcam: - return Webcam( - camera_index=0, - fps=15, - stereo_slice="left", - camera_info=zed.CameraInfo.SingleWebcam, - ) - - _camera = ( autoconnect( CameraModule.blueprint( @@ -115,7 +106,12 @@ def _create_webcam() -> Webcam: frame_id="sensor", child_frame_id="camera_link", ), - hardware=_create_webcam, + hardware=WebcamConfig( + camera_index=0, + fps=15, + stereo_slice="left", + camera_info=zed.CameraInfo.SingleWebcam, + ), ), ) if not global_config.simulation diff --git a/dimos/robot/unitree/g1/test_g1_teleop.py b/dimos/robot/unitree/g1/test_g1_teleop.py index b264280339..5363d3117b 100644 --- a/dimos/robot/unitree/g1/test_g1_teleop.py +++ b/dimos/robot/unitree/g1/test_g1_teleop.py @@ -31,6 +31,10 @@ from dimos.core.coordination.blueprints import Blueprint from dimos.manipulation.planning.spec.validation import prepare_robot_model from dimos.manipulation.visualization.viser.config import ViserVisualizationConfig +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.imitation_msgs.EpisodeStatus import EpisodeStatus +from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.sensor_msgs.JointState import JointState from dimos.robot.unitree.g1.blueprints.basic.unitree_g1_groot_wbc import ( _G1_TELEOP_MODEL, _G1GrootCoordinator, @@ -141,16 +145,23 @@ def test_g1_teleop_excludes_navigation_and_legacy_visualization() -> None: ) -def test_g1_collection_streams_do_not_require_world_poses() -> None: - recorder_kwargs = _module_kwargs(unitree_g1_teleop, G1CollectionRecorder) - - assert recorder_kwargs["poseless_streams"] == [ - "color_image", - "status", - "left_cartesian_command", - "right_cartesian_command", - "coordinator_joint_state", - ] +async def test_g1_collection_accepts_blueprint_config_and_records_without_world_poses(tmp_path): + kwargs = _module_kwargs(unitree_g1_teleop, G1CollectionRecorder) + recorder = G1CollectionRecorder(**{**kwargs, "db_path": str(tmp_path / "session.db")}) + try: + ports = recorder._data_ports() + assert {name: port.type for name, port in ports.items()} == { + "color_image": Image, + "coordinator_joint_state": JointState, + "status": EpisodeStatus, + "left_cartesian_command": PoseStamped, + "right_cartesian_command": PoseStamped, + } + assert recorder.config.record_tf is False + for name in ports: + assert await recorder._resolve_pose(name, JointState(ts=1.0), 1.0) is None + finally: + recorder.stop() @pytest.mark.self_hosted diff --git a/dimos/teleop/webxr/module.py b/dimos/teleop/webxr/module.py index e6925b93d2..a48c234609 100644 --- a/dimos/teleop/webxr/module.py +++ b/dimos/teleop/webxr/module.py @@ -43,8 +43,8 @@ from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In, Out -from dimos.imitation.collection.episode_monitor import EpisodeStatus from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.imitation_msgs.EpisodeStatus import EpisodeStatus from dimos.msgs.sensor_msgs.Joy import Joy from dimos.teleop.utils.teleop_transforms import webxr_to_robot from dimos.teleop.webxr.body_tracking import BodyTrackingMode, BodyTrackingSnapshot diff --git a/dimos/teleop/webxr/test_module.py b/dimos/teleop/webxr/test_module.py index 06c415a217..2a33781cf6 100644 --- a/dimos/teleop/webxr/test_module.py +++ b/dimos/teleop/webxr/test_module.py @@ -24,8 +24,8 @@ import pytest import pytest_mock -from dimos.imitation.collection.episode_monitor import EpisodeStatus from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.imitation_msgs.EpisodeStatus import EpisodeStatus from dimos.msgs.sensor_msgs.Joy import Joy from dimos.teleop.webxr.body_tracking import BodyTrackingSnapshot from dimos.teleop.webxr.controller_types import ( diff --git a/examples/native-modules/rust/Cargo.toml b/examples/native-modules/rust/Cargo.toml index d2c1d48d48..9ee0334293 100644 --- a/examples/native-modules/rust/Cargo.toml +++ b/examples/native-modules/rust/Cargo.toml @@ -21,7 +21,7 @@ path = "src/tf_broadcaster.rs" [dependencies] dimos-module = { path = "../../../native/rust/dimos-module" } -lcm-msgs = { git = "https://github.com/dimensionalOS/dimos-lcm.git", branch = "rust-codegen" } +lcm-msgs = { git = "https://github.com/dimensionalOS/dimos-lcm.git", rev = "dd2159513ebfaa7ebc5fc32bf60209cd09aa1ca5" } tokio = { version = "1", features = ["rt-multi-thread", "macros", "time"] } serde = { version = "1", features = ["derive"] } tracing = "0.1" diff --git a/native/python/lerobot/dimos_lerobot/dataprep.py b/native/python/lerobot/dimos_lerobot/dataprep.py new file mode 100644 index 0000000000..a3b3fff253 --- /dev/null +++ b/native/python/lerobot/dimos_lerobot/dataprep.py @@ -0,0 +1,218 @@ +# 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 LeRobot dataset writer and isolated command entry point.""" + +from __future__ import annotations + +from collections.abc import Iterator +from contextlib import redirect_stdout, suppress +from pathlib import Path +import sys +from typing import Any, Protocol, cast + +from lerobot.datasets import LeRobotDatasetMetadata +from lerobot.datasets.lerobot_dataset import LeRobotDataset +import numpy as np +from numpy.typing import NDArray + +from dimos.imitation.dataprep._lerobot_protocol import ( + REQUEST_ADAPTER, + BuildRequest, + BuildResult, + InspectRequest, + InspectResult, +) +from dimos.imitation.dataprep.build import run_dataprep +from dimos.imitation.dataprep.core import OutputConfig, Sample, summarize_lengths + + +class _WritableDataset(Protocol): + root: Path + + def add_frame(self, frame: dict[str, Any]) -> None: ... + + def save_episode(self, *, parallel_encoding: bool = True) -> None: ... + + def clear_episode_buffer(self) -> None: ... + + def finalize(self) -> None: ... + + +def _task(sample: Sample) -> str: + value = sample.task_label + if not isinstance(value, str) or not value.strip(): + raise ValueError("every LeRobot frame requires an episode task label") + return value + + +def _sample_features(sample: Sample) -> dict[str, NDArray[Any]]: + return { + **sample.observation, + **sample.action, + **{f"complementary_info.{key}": value for key, value in sample.complementary_info.items()}, + } + + +def write(samples: Iterator[Sample], output: OutputConfig) -> Path: + """Write synchronized DimOS samples through LeRobot's native dataset API.""" + repo_id = output.metadata.get("repo_id") + if not isinstance(repo_id, str) or not repo_id.strip(): + raise ValueError("LeRobot output.metadata.repo_id is required") + fps_value = output.metadata.get("fps") + if ( + not isinstance(fps_value, (int, float)) + or isinstance(fps_value, bool) + or fps_value <= 0 + or not float(fps_value).is_integer() + ): + raise ValueError("LeRobot output.metadata.fps must be a positive integer") + fps = int(fps_value) + + iterator = iter(samples) + try: + first = next(iterator) + except StopIteration as error: + raise ValueError("cannot create a LeRobot dataset without samples") from error + + raw_schema = output.metadata.get("feature_schema") + if not isinstance(raw_schema, dict) or not raw_schema: + raise ValueError("LeRobot output.metadata.feature_schema is required") + features: dict[str, dict[str, Any]] = {} + for key, raw in raw_schema.items(): + if not isinstance(key, str) or not isinstance(raw, dict): + raise ValueError("feature_schema must map feature names to definitions") + dtype = raw.get("dtype") + shape = raw.get("shape") + names = raw.get("names") + if not isinstance(dtype, str): + raise ValueError(f"feature {key!r} requires a dtype string") + if not isinstance(shape, (list, tuple)) or not all( + isinstance(value, int) and value > 0 for value in shape + ): + raise ValueError(f"feature {key!r} requires a positive integer shape") + if not isinstance(names, list) or not all(isinstance(name, str) for name in names): + raise ValueError(f"feature {key!r} requires string axis names") + features[key] = {"dtype": dtype, "shape": tuple(shape), "names": names} + + first_values = _sample_features(first) + if set(first_values) != set(features): + raise ValueError( + f"sample features {sorted(first_values)} do not match schema {sorted(features)}" + ) + dataset = cast( + "_WritableDataset", + LeRobotDataset.create( + repo_id=repo_id, + fps=fps, + features=features, + root=output.path, + robot_type=output.metadata.get("robot_type"), + use_videos=True, + ), + ) + current_episode: str | None = None + finished: set[str] = set() + + def add(sample: Sample) -> None: + nonlocal current_episode + if current_episode is not None and sample.episode_id != current_episode: + dataset.save_episode(parallel_encoding=False) + finished.add(current_episode) + if sample.episode_id in finished: + raise ValueError(f"episode {sample.episode_id!r} is not contiguous") + current_episode = sample.episode_id + values = _sample_features(sample) + if set(values) != set(features): + raise ValueError(f"sample feature keys changed in episode {sample.episode_id}") + frame: dict[str, Any] = {"task": _task(sample)} + for key, definition in features.items(): + value = np.asarray(values[key]) + shape = tuple(definition["shape"]) + if value.shape != shape: + raise ValueError(f"{key} shape changed from {shape} to {value.shape}") + dtype = definition["dtype"] + frame[key] = ( + value.astype(np.uint8, copy=False) + if dtype == "video" + else value.astype(np.dtype(dtype), copy=False) + ) + dataset.add_frame(frame) + + try: + add(first) + for sample in iterator: + add(sample) + dataset.save_episode(parallel_encoding=False) + dataset.finalize() + except BaseException: + with suppress(Exception): + dataset.clear_episode_buffer() + with suppress(Exception): + dataset.finalize() + raise + return Path(dataset.root) + + +def inspect_dataset(path: Path) -> dict[str, Any]: + """Summarize a local dataset through LeRobot's public metadata interface.""" + metadata = LeRobotDatasetMetadata(repo_id="local/dataset", root=path) + observation: dict[str, Any] = {} + action: dict[str, Any] = {} + metadata_columns = {"timestamp", "frame_index", "episode_index", "index", "task_index"} + for name, feature in metadata.features.items(): + if name in metadata_columns: + continue + entry = {"shape": feature.get("shape"), "dtype": feature.get("dtype")} + if name.startswith("observation"): + observation[name] = entry + elif name.startswith("action"): + action[name] = entry + lengths = list(metadata.episodes["length"]) if metadata.episodes is not None else [] + return { + "format": "lerobot", + "version": metadata.info.codebase_version, + "path": str(metadata.root), + "episodes": metadata.total_episodes, + "frames": metadata.total_frames, + "fps": metadata.fps, + "robot": metadata.robot_type, + "observation": observation, + "action": action, + "episode_lengths": summarize_lengths(lengths), + "shapes_uniform": True, + "has_stats": metadata.stats is not None, + } + + +def main(argv: list[str] | None = None) -> None: + """Execute one typed build or inspect request from stdin.""" + args = list(sys.argv[1:] if argv is None else argv) + if args: + raise SystemExit("usage: python -m dimos_lerobot.dataprep") + request = REQUEST_ADAPTER.validate_json(sys.stdin.read()) + protocol_stdout = sys.stdout + result: BuildResult | InspectResult + with redirect_stdout(sys.stderr): + if isinstance(request, BuildRequest): + result = BuildResult(path=run_dataprep(request.config, writer=write)) + elif isinstance(request, InspectRequest): + result = InspectResult(info=inspect_dataset(request.path)) + else: + raise TypeError(f"unsupported request {type(request).__name__}") + protocol_stdout.write(result.model_dump_json()) + + +if __name__ == "__main__": + main() diff --git a/native/python/lerobot/dimos_lerobot/dataprep_tests.py b/native/python/lerobot/dimos_lerobot/dataprep_tests.py new file mode 100644 index 0000000000..1aa22cd95c --- /dev/null +++ b/native/python/lerobot/dimos_lerobot/dataprep_tests.py @@ -0,0 +1,147 @@ +# 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. + +from collections.abc import Iterator +from io import StringIO +import json +from pathlib import Path +import sys + +from dimos_lerobot import dataprep +from dimos_lerobot.dataprep import write +import numpy as np +import pytest +import pytest_mock + +from dimos.imitation.dataprep._lerobot_protocol import BuildRequest, BuildResult +from dimos.imitation.dataprep.core import DataPrepConfig, OutputConfig, Sample + +JOINTS = [f"arm/joint{index}" for index in range(1, 7)] + ["arm/gripper"] + + +def test_module_entry_point_executes_typed_build_request( + tmp_path: Path, + mocker: pytest_mock.MockerFixture, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + config = DataPrepConfig( + source="recording.mcap", + output=OutputConfig(format="lerobot", path=tmp_path / "dataset"), + ) + request = BuildRequest(config=config) + run = mocker.patch.object(dataprep, "run_dataprep", return_value=tmp_path / "dataset") + monkeypatch.setattr(sys, "stdin", StringIO(request.model_dump_json())) + + dataprep.main([]) + + result = BuildResult.model_validate_json(capsys.readouterr().out) + assert result.path == tmp_path / "dataset" + assert run.call_args.args[0].source == "recording.mcap" + assert run.call_args.kwargs["writer"] is write + + +@pytest.mark.parametrize("args", [["config.json"], ["one", "two"]]) +def test_module_entry_point_rejects_arguments(args: list[str]) -> None: + with pytest.raises(SystemExit, match="usage: python -m dimos_lerobot.dataprep"): + dataprep.main(args) + + +def samples() -> Iterator[Sample]: + for episode, task in (("first", "pick"), ("second", "place")): + for frame in range(3): + value = float(frame + (10 if episode == "second" else 0)) + yield Sample( + ts=value, + episode_id=episode, + observation={ + "observation.images.wrist": np.full((64, 64, 3), frame, dtype=np.uint8), + "observation.state": np.full(7, value, dtype=np.float32), + "observation.effort": np.full(7, value * 0.1, dtype=np.float32), + }, + action={"action": np.full(7, value + 1, dtype=np.float32)}, + task_label=task, + complementary_info={"is_filled": np.asarray([False])}, + ) + + +def output(path: Path) -> OutputConfig: + return OutputConfig( + format="lerobot", + path=path, + metadata={ + "repo_id": "local/openyam-test", + "fps": 30, + "robot_type": "openyam", + "feature_schema": { + "observation.images.wrist": { + "dtype": "video", + "shape": [64, 64, 3], + "names": ["height", "width", "channels"], + }, + "observation.state": {"dtype": "float32", "shape": [7], "names": JOINTS}, + "observation.effort": {"dtype": "float32", "shape": [7], "names": JOINTS}, + "action": {"dtype": "float32", "shape": [7], "names": JOINTS}, + "complementary_info.is_filled": { + "dtype": "bool", + "shape": [1], + "names": ["is_filled"], + }, + }, + }, + ) + + +def test_native_writer_creates_canonical_openyam_dataset(tmp_path: Path) -> None: + root = write(samples(), output(tmp_path / "dataset")) + + info = json.loads((root / "meta" / "info.json").read_text()) + assert info["total_episodes"] == 2 + assert info["total_frames"] == 6 + assert info["fps"] == 30 + assert info["robot_type"] == "openyam" + assert set(info["features"]) >= { + "observation.images.wrist", + "observation.state", + "action", + "observation.effort", + "complementary_info.is_filled", + } + assert info["features"]["observation.state"]["names"] == JOINTS + + summary = dataprep.inspect_dataset(root) + assert summary["episodes"] == 2 + assert summary["frames"] == 6 + assert summary["episode_lengths"] == { + "min": 3, + "max": 3, + "mean": 3.0, + "uniform": True, + } + assert summary["observation"]["observation.images.wrist"]["dtype"] == "video" + + +def test_native_writer_requires_repo_id(tmp_path: Path) -> None: + config = output(tmp_path / "dataset").model_copy(update={"metadata": {"fps": 30}}) + + with pytest.raises(ValueError, match="repo_id is required"): + write(samples(), config) + + +def test_native_writer_rejects_fractional_fps(tmp_path: Path) -> None: + config = output(tmp_path / "dataset") + config.metadata["fps"] = 14.5 + + with pytest.raises(ValueError, match="positive integer"): + write(samples(), config) diff --git a/native/python/lerobot/dimos_lerobot/mcap_dataprep_tests.py b/native/python/lerobot/dimos_lerobot/mcap_dataprep_tests.py new file mode 100644 index 0000000000..59b60a81cc --- /dev/null +++ b/native/python/lerobot/dimos_lerobot/mcap_dataprep_tests.py @@ -0,0 +1,235 @@ +# 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. + +"""Full native-MCAP to LeRobot conversion in the locked policy environment.""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any, Literal + +from dimos_lerobot.dataprep import write +from lerobot.datasets.lerobot_dataset import LeRobotDataset +from mcap.writer import Writer as McapWriter +import numpy as np +import pytest + +from dimos.imitation.collection.profile import CollectionFeature, CollectionProfile +from dimos.imitation.collection.recording import RecordingSchema +from dimos.imitation.dataprep.build import run_dataprep +from dimos.imitation.dataprep.core import ( + OutputConfig, + SyncConfig, +) +from dimos.memory.codecs.jpeg import JpegCodec +from dimos.msgs.imitation_msgs.EpisodeStatus import EpisodeStatus +from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.msgs.sensor_msgs.JointState import JointState + +JOINTS = [f"joint_{i}" for i in range(14)] + + +def _channel( + writer: McapWriter, + name: str, + payload_type: type[Any], + message_encoding: str = "lcm", +) -> int: + return int( + writer.register_channel( + topic=name, + message_encoding=message_encoding, + schema_id=0, + metadata={ + "dimos.payload_type": f"{payload_type.__module__}.{payload_type.__qualname__}", + "dimos.observation_time": "publish_time", + }, + ) + ) + + +def _add(writer: McapWriter, channel_id: int, ts: float, message: Any) -> None: + timestamp_ns = round(ts * 1_000_000_000) + writer.add_message( + channel_id=channel_id, + log_time=timestamp_ns, + publish_time=timestamp_ns, + data=message if isinstance(message, bytes) else message.lcm_encode(), + ) + + +def _record(path: Path, camera_count: int) -> None: + with path.open("wb") as output: + writer = McapWriter(output) + writer.start(profile="dimos", library="test") + image_channels = [_channel(writer, f"view_{i}", Image, "jpeg") for i in range(camera_count)] + state_channel = _channel(writer, "measured", JointState) + action_channel = _channel(writer, "commanded", JointState) + status_channel = _channel(writer, "status", EpisodeStatus) + events: tuple[Literal["save", "discard"], ...] = ("save", "discard") + for episode, event in enumerate(events): + start_time = 20.0 + episode + _add( + writer, + status_channel, + start_time, + EpisodeStatus( + ts=start_time, + state="recording", + last_event="start", + episodes_saved=episode, + episodes_discarded=0, + task_label="pick", + ), + ) + for frame in range(3): + ts = start_time + frame / 30.0 + for image_channel in image_channels: + _add( + writer, + image_channel, + ts, + JpegCodec().encode( + Image( + ts=ts, + frame_id="wrist_camera_link", + format=ImageFormat.RGB, + data=np.full((64, 64, 3), frame, dtype=np.uint8), + ) + ), + ) + _add( + writer, + state_channel, + ts, + JointState( + ts=ts, + frame_id="coordinator", + name=list(reversed(JOINTS)), + position=[float(frame + i) for i in reversed(range(len(JOINTS)))], + velocity=[0.0] * len(JOINTS), + effort=[0.0] * len(JOINTS), + ), + ) + _add( + writer, + action_channel, + ts, + JointState( + ts=ts, + frame_id="coordinator", + name=list(reversed(JOINTS)), + position=[float(frame + i) + 0.5 for i in reversed(range(len(JOINTS)))], + velocity=[], + effort=[], + ), + ) + end = start_time + 2 / 30.0 + _add( + writer, + status_channel, + end, + EpisodeStatus( + ts=end, + state="idle", + last_event=event, + episodes_saved=1, + episodes_discarded=0, + task_label="pick", + ), + ) + writer.finish() + + +@pytest.mark.parametrize("camera_count", [1, 2, 4]) +def test_mcap_converts_to_lerobot_dataset(tmp_path: Path, camera_count: int) -> None: + source = tmp_path / "session.mcap" + destination = tmp_path / "dataset" + _record(source, camera_count) + profile = CollectionProfile( + name="dual-test", + robot_type="dual_openyam", + observations={ + **{ + f"observation.images.view_{i}": CollectionFeature( + stream=f"view_{i}", + message_type=Image, + field="data", + dtype="video", + shape=(64, 64, 3), + names=["height", "width", "channels"], + ) + for i in range(camera_count) + }, + "observation.state": CollectionFeature( + stream="measured", + message_type=JointState, + field="position", + dtype="float32", + shape=(len(JOINTS),), + names=JOINTS, + ), + "observation.velocity": CollectionFeature( + stream="measured", + message_type=JointState, + field="velocity", + dtype="float32", + shape=(len(JOINTS),), + names=JOINTS, + ), + }, + actions={ + "action": CollectionFeature( + stream="commanded", + message_type=JointState, + field="position", + dtype="float32", + shape=(len(JOINTS),), + names=JOINTS, + ), + }, + sync=SyncConfig(anchor="observation.images.view_0", rate_hz=30, tolerance_ms=20), + ) + directory = tmp_path / "session" + directory.mkdir() + source.rename(directory / "recording.mcap") + (directory / "schema.json").write_text(profile.to_schema().model_dump_json()) + moved = tmp_path / "moved" + directory.rename(moved) + config = RecordingSchema.read(moved).dataprep_config( + moved, + OutputConfig( + format="lerobot", + path=destination, + metadata={"repo_id": "local/openyam-mcap"}, + ), + ) + # Exercise the exact JSON boundary used by the isolated converter. + config = type(config).model_validate_json(config.model_dump_json()) + root = run_dataprep(config, writer=write) + + info = json.loads((root / "meta" / "info.json").read_text()) + assert info["total_episodes"] == 1 + assert info["total_frames"] == 3 + assert info["fps"] == 30 + assert info["features"]["action"]["names"] == JOINTS + + dataset = LeRobotDataset("local/openyam-mcap", root=root) + frame = dataset[0] + for i in range(camera_count): + assert tuple(frame[f"observation.images.view_{i}"].shape) == (3, 64, 64) + np.testing.assert_allclose(frame["observation.state"].numpy(), np.arange(len(JOINTS))) + np.testing.assert_allclose(frame["action"].numpy(), np.arange(len(JOINTS)) + 0.5) + np.testing.assert_allclose(frame["observation.velocity"].numpy(), np.zeros(len(JOINTS))) diff --git a/native/python/lerobot/pyproject.toml b/native/python/lerobot/pyproject.toml index 81a4e267b6..ed3d582172 100644 --- a/native/python/lerobot/pyproject.toml +++ b/native/python/lerobot/pyproject.toml @@ -7,7 +7,8 @@ name = "dimos-lerobot-runtime" version = "0.1.0" requires-python = ">=3.12,<3.13" dependencies = [ - "lerobot==0.6.0", + "lerobot[dataset,viz]==0.6.0", + "mcap>=1.2", "opencv-contrib-python>=4.8,<5", "transformers[torch]>=5.4,<5.6", ] @@ -25,6 +26,9 @@ override-dependencies = [ # LeRobot's headless OpenCV wheel owns the same cv2/ tree as DimOS's # opencv-contrib-python dependency. Contrib is the required superset. "opencv-python-headless; sys_platform == 'never'", + # LeRobot selects TorchCodec whenever it is importable, but its native wheel + # cannot load against newer system FFmpeg ABIs. Use LeRobot's PyAV backend. + "torchcodec; sys_platform == 'never'", ] [tool.setuptools.packages.find] @@ -46,6 +50,7 @@ mypy_path = "../../.." [[tool.mypy.overrides]] module = [ "lerobot.configs.policies", + "lerobot.datasets", "lerobot.datasets.lerobot_dataset", "lerobot.policies.factory", "lerobot.policies.pretrained", @@ -56,3 +61,7 @@ module = [ ] follow_untyped_imports = true ignore_missing_imports = true + +[[tool.mypy.overrides]] +module = ["mcap", "mcap.*"] +follow_imports = "skip" diff --git a/native/python/lerobot/uv.lock b/native/python/lerobot/uv.lock index e0fec2968d..66fe2d4b74 100644 --- a/native/python/lerobot/uv.lock +++ b/native/python/lerobot/uv.lock @@ -10,7 +10,10 @@ resolution-markers = [ ] [manifest] -overrides = [{ name = "opencv-python-headless", marker = "sys_platform == 'never'" }] +overrides = [ + { name = "opencv-python-headless", marker = "sys_platform == 'never'" }, + { name = "torchcodec", marker = "sys_platform == 'never'" }, +] [[package]] name = "accelerate" @@ -30,6 +33,64 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a8/db/253133d7e7cb40d3af384bb2f5c0b4a2b7fdcffbc95c688cc67a20a3c103/accelerate-1.14.0-py3-none-any.whl", hash = "sha256:e94390c2863b873be18f623f9df48a0d8fe5eff13ea7f1a00092b0a7904888c6", size = 389246, upload-time = "2026-06-11T13:45:50.477Z" }, ] +[[package]] +name = "aiohappyeyeballs" +version = "2.7.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/ce/f4/eec0465c2f67b2664688d0240b3212d5196fd89e741df67ddb81f8d35658/aiohappyeyeballs-2.7.1.tar.gz", hash = "sha256:065665c041c42a5938ed220bdcd7230f22527fbec085e1853d2402c8a3615d9d", size = 24757, upload-time = "2026-07-01T17:11:55.501Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/71/43/1947f06babed6b3f1d7f38b0c767f52df66bfb2bc10b468c4a7de9eceff2/aiohappyeyeballs-2.7.1-py3-none-any.whl", hash = "sha256:9243213661e29250eb41368e5daa826fc017156c3b8a11440826b2e3ed376472", size = 15038, upload-time = "2026-07-01T17:11:54.055Z" }, +] + +[[package]] +name = "aiohttp" +version = "3.14.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "aiohappyeyeballs" }, + { name = "aiosignal" }, + { name = "attrs" }, + { name = "frozenlist" }, + { name = "multidict" }, + { name = "propcache" }, + { name = "typing-extensions" }, + { name = "yarl" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/58/d9/22ce5786ac0c1653ae8b6c23bded02c1686d11f0dbb45b31ce128e0df985/aiohttp-3.14.3.tar.gz", hash = "sha256:9491196535a88924a60afd5b5f434b5b203b6cc616250878dbdb223a8f7844bc", size = 7971213, upload-time = "2026-07-23T01:57:27.037Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/18/d4/eb96299230e20acf2efae207cb8d69051f1f68e357e5ea5e479bf6fb097a/aiohttp-3.14.3-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:39aded8c7f3b935b54aab1d8d73c70ec0ee2d3ec3b943e0e86611bc150ba47f5", size = 754690, upload-time = "2026-07-23T01:53:47.332Z" }, + { url = "https://files.pythonhosted.org/packages/88/11/e7a70a209eb9a067c0d3212b518a0134e3484f5178c7533878b6b514d469/aiohttp-3.14.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:5bcb6ff3fdab1258a192679ff1a05d44f59626430aa05cd1a9d2447423599228", size = 509484, upload-time = "2026-07-23T01:53:51.159Z" }, + { url = "https://files.pythonhosted.org/packages/30/07/4bbc222cc8dbe31d4c3e8a5baad2286e4d42026ac0c570027b89afce6344/aiohttp-3.14.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:617105e2c3018ee38d0c8ce5ee3c84f621a6d8b9f723202aacaff28449ca91ee", size = 511949, upload-time = "2026-07-23T01:53:55.083Z" }, + { url = "https://files.pythonhosted.org/packages/54/b9/42e74c46b7b7c794b995bbc1f573fb48950c38b19d8600c62a6804ee2d67/aiohttp-3.14.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f631fe87a6f30df5fbe6d79640b25e4cffb38c31c7fb6f10871517b84b0f8c1a", size = 1765282, upload-time = "2026-07-23T01:53:59.662Z" }, + { url = "https://files.pythonhosted.org/packages/6b/ed/62bc4d74363ad346d518e0720363a949f63e2e23439a79eb5813d4d29bb3/aiohttp-3.14.3-cp312-cp312-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:a94dbaae5ae27bd849c93570669bff91e0510f33a80805738e3de72a7be0447b", size = 1741511, upload-time = "2026-07-23T01:54:04.063Z" }, + { url = "https://files.pythonhosted.org/packages/d0/9f/181e8a8bc79e47d13c7fc4540bd7a3b729d9505609c61f392a8dd2fbfe55/aiohttp-3.14.3-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:8f2f1c4c032c7cedd7d8da6f54c97b70266c6570c3108d3fdffee7188bb70529", size = 1810680, upload-time = "2026-07-23T01:54:09.882Z" }, + { url = "https://files.pythonhosted.org/packages/5c/9a/dec94d6ad694552fe3424e3f1928d7a606a5d9d9433a04e7ecdd9d38ae7f/aiohttp-3.14.3-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:ea05e1f97ceea523942d9b2a7d7c0359d781d683d6b043f5943a602b14da4787", size = 1905646, upload-time = "2026-07-23T01:54:13.475Z" }, + { url = "https://files.pythonhosted.org/packages/52/b7/7cd31f29d6055bd711ae6e669367fba6f5ae9de463910a793e30556a8db7/aiohttp-3.14.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:543906c127fb1d929b95076db19b83fa2d46751006ff1e23b093aa5ac4d8db42", size = 1792122, upload-time = "2026-07-23T01:54:15.752Z" }, + { url = "https://files.pythonhosted.org/packages/66/73/10b1ef93afa61f4963c746257b70ced619cf31a4798671de5fdb2608501d/aiohttp-3.14.3-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:0a5ff2dfbb9ce645fa5b8ef3e02c6c0b9cc3f6030ff863d0c51fffc50cb5541b", size = 1591127, upload-time = "2026-07-23T01:54:19.489Z" }, + { url = "https://files.pythonhosted.org/packages/49/ed/3b203fa6de1b338c14acdc06bf6ca9b043b7944f005966958c2ced932cde/aiohttp-3.14.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:041badb8f84396357c4d3ad26de6afd7a32b112f43d3c63045c0c8278cfd2043", size = 1725210, upload-time = "2026-07-23T01:54:24.129Z" }, + { url = "https://files.pythonhosted.org/packages/28/b7/1c2aab8c706436dcc28598452488ac9cd7c409da815237c28c27d58993e6/aiohttp-3.14.3-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:530125ee1163c4219af35dc3aa1206e541e7b31b6efc1a3f93b70a136f65d427", size = 1764848, upload-time = "2026-07-23T01:54:27.973Z" }, + { url = "https://files.pythonhosted.org/packages/54/50/94c28f08b131c4bf10984ea2c7a536c9920608bb2d6e7f95642c30cc87b7/aiohttp-3.14.3-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:c8653fd547c93a61aadc612007790f5555cdd18946fa48cf45e26d8ea4ea473d", size = 1777102, upload-time = "2026-07-23T01:54:31.775Z" }, + { url = "https://files.pythonhosted.org/packages/13/d4/e7d09ba7d345fb2d74440fd2fa033c5e079fac05552927705986f41a364f/aiohttp-3.14.3-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:89176250f686cb9853c0fb7ead90e639e915b84a6f43eedc2a4e7ec21f1037f0", size = 1580205, upload-time = "2026-07-23T01:54:34.518Z" }, + { url = "https://files.pythonhosted.org/packages/a3/84/072a91d68e1e1eb587985b54baab94221277f877e8ef274fc213a0ceae28/aiohttp-3.14.3-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:3a26434dafe408229ff3403458ca58de24fb51936504decac49ce6755f77e59d", size = 1797219, upload-time = "2026-07-23T01:54:36.995Z" }, + { url = "https://files.pythonhosted.org/packages/e0/eb/aad34e897e668424d6e995da5dff8a4a09af93363d3392488772957a63aa/aiohttp-3.14.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:d1558173930a5a8d3069cee5c92fc91c87c4dbcb099debbb3622053717145a19", size = 1768629, upload-time = "2026-07-23T01:54:40.103Z" }, + { url = "https://files.pythonhosted.org/packages/b6/2b/6bb88ddba0fecd9122aa3ebcad25996cf6c083a4a7040dbb3a4f97972af6/aiohttp-3.14.3-cp312-cp312-win32.whl", hash = "sha256:16100ad3ab8d649fdfbee87602d9d2dcdca9df0b9eda8a1b5fdc0d41f96da559", size = 451481, upload-time = "2026-07-23T01:54:42.547Z" }, + { url = "https://files.pythonhosted.org/packages/76/9b/f2f8f108da17ecef2cc3efc424e8b7ad3782b1a8360f7b8eae8ced84f6ea/aiohttp-3.14.3-cp312-cp312-win_amd64.whl", hash = "sha256:33a2d7c28d33797a2e99923dffa63f83d908a19b6bf26cfe80fa790aa5e1a75a", size = 476845, upload-time = "2026-07-23T01:54:44.853Z" }, + { url = "https://files.pythonhosted.org/packages/3e/44/28dac80a8941b604f4da10ce21097614ca1bf905ce93dca28d8d7de9c1e7/aiohttp-3.14.3-cp312-cp312-win_arm64.whl", hash = "sha256:362a3fd481769cac1a824514bcd86fda51c65e8fe6e051099e008fddde6db17c", size = 448050, upload-time = "2026-07-23T01:54:47.087Z" }, +] + +[[package]] +name = "aiosignal" +version = "1.4.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "frozenlist" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/61/62/06741b579156360248d1ec624842ad0edf697050bbaf7c3e46394e106ad1/aiosignal-1.4.0.tar.gz", hash = "sha256:f47eecd9468083c2029cc99945502cb7708b082c232f9aca65da147157b251c7", size = 25007, upload-time = "2025-07-03T22:54:43.528Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fb/76/641ae371508676492379f16e2fa48f4e2c11741bd63c48be4b12a6b09cba/aiosignal-1.4.0-py3-none-any.whl", hash = "sha256:053243f8b92b990551949e63930a839ff0cf0b0ebbe0597b0f3fb19e1a0fe82e", size = 7490, upload-time = "2025-07-03T22:54:42.156Z" }, +] + [[package]] name = "annotated-doc" version = "0.0.5" @@ -52,6 +113,30 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/da/35/f2287558c17e29fafc8ef3daf819bb9834061cfa43bff8014f7df7f63bdc/anyio-4.14.2-py3-none-any.whl", hash = "sha256:9f505dda5ac9f0c8309b5e8bd445a8c2bf7246f3ce950121e45ea15bc41d1494", size = 125813, upload-time = "2026-07-12T20:29:05.763Z" }, ] +[[package]] +name = "attrs" +version = "26.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/9a/8e/82a0fe20a541c03148528be8cac2408564a6c9a0cc7e9171802bc1d26985/attrs-26.1.0.tar.gz", hash = "sha256:d03ceb89cb322a8fd706d4fb91940737b6642aa36998fe130a9bc96c985eff32", size = 952055, upload-time = "2026-03-19T14:22:25.026Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/64/b4/17d4b0b2a2dc85a6df63d1157e028ed19f90d4cd97c36717afef2bc2f395/attrs-26.1.0-py3-none-any.whl", hash = "sha256:c647aa4a12dfbad9333ca4e71fe62ddc36f4e63b2d260a37a8b83d2f043ac309", size = 67548, upload-time = "2026-03-19T14:22:23.645Z" }, +] + +[[package]] +name = "av" +version = "15.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e9/c3/83e6e73d1592bc54436eae0bc61704ae0cff0c3cfbde7b58af9ed67ebb49/av-15.1.0.tar.gz", hash = "sha256:39cda2dc810e11c1938f8cb5759c41d6b630550236b3365790e67a313660ec85", size = 3774192, upload-time = "2025-08-30T04:41:56.076Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d1/58/de78b276d20db6ffcd4371283df771721a833ba525a3d57e753d00a9fe79/av-15.1.0-cp312-cp312-macosx_13_0_arm64.whl", hash = "sha256:40c5df37f4c354ab8190c6fd68dab7881d112f527906f64ca73da4c252a58cee", size = 21760991, upload-time = "2025-08-30T04:40:00.801Z" }, + { url = "https://files.pythonhosted.org/packages/56/cc/45f85775304ae60b66976360d82ba5b152ad3fd91f9267d5020a51e9a828/av-15.1.0-cp312-cp312-macosx_13_0_x86_64.whl", hash = "sha256:af455ce65ada3d361f80c90c810d9bced4db5655ab9aa513024d6c71c5c476d5", size = 26953097, upload-time = "2025-08-30T04:40:03.998Z" }, + { url = "https://files.pythonhosted.org/packages/f3/f8/2d781e5e71d02fc829487e775ccb1185e72f95340d05f2e84eb57a11e093/av-15.1.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:86226d2474c80c3393fa07a9c366106029ae500716098b72b3ec3f67205524c3", size = 38319710, upload-time = "2025-08-30T04:40:07.701Z" }, + { url = "https://files.pythonhosted.org/packages/ac/13/37737ef2193e83862ccacff23580c39de251da456a1bf0459e762cca273c/av-15.1.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:11326f197e7001c4ca53a83b2dbc67fd39ddff8cdf62ce6be3b22d9f3f9338bd", size = 39915519, upload-time = "2025-08-30T04:40:11.066Z" }, + { url = "https://files.pythonhosted.org/packages/26/e9/e8032c7b8f2a4129a03f63f896544f8b7cf068e2db2950326fa2400d5c47/av-15.1.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:a631ea879cc553080ee62874f4284765c42ba08ee0279851a98a85e2ceb3cc8d", size = 40286166, upload-time = "2025-08-30T04:40:14.561Z" }, + { url = "https://files.pythonhosted.org/packages/e2/23/612c0fd809444d04b8387a2dfd942ccc77829507bd78a387ff65a9d98c24/av-15.1.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:8f383949b010c3e731c245f80351d19dc0c08f345e194fc46becb1cb279be3ff", size = 41150592, upload-time = "2025-08-30T04:40:17.951Z" }, + { url = "https://files.pythonhosted.org/packages/15/74/6f8e38a3b0aea5f28e72813672ff45b64615f2c69e6a4a558718c95edb9f/av-15.1.0-cp312-cp312-win_amd64.whl", hash = "sha256:d5921aa45f4c1f8c1a8d8185eb347e02aa4c3071278a2e2dd56368d54433d643", size = 31336093, upload-time = "2025-08-30T04:40:21.393Z" }, +] + [[package]] name = "certifi" version = "2026.7.22" @@ -221,12 +306,47 @@ nvtx = [ { name = "nvidia-nvtx", marker = "sys_platform == 'linux'" }, ] +[[package]] +name = "datasets" +version = "4.8.5" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "dill" }, + { name = "filelock" }, + { name = "fsspec", extra = ["http"] }, + { name = "httpx" }, + { name = "huggingface-hub" }, + { name = "multiprocess" }, + { name = "numpy" }, + { name = "packaging" }, + { name = "pandas" }, + { name = "pyarrow" }, + { name = "pyyaml" }, + { name = "requests" }, + { name = "tqdm" }, + { name = "xxhash" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/66/34/14cd8e76f907f7d4dca2334cfeec9f81d30fd15c25a015f99aaea694eaed/datasets-4.8.5.tar.gz", hash = "sha256:0f0c1c3d56ffff2c93b2f4c63c95bac94f3d7e8621aea2a2a576275233bba772", size = 605649, upload-time = "2026-04-27T15:43:57.384Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/65/99/00f3196036501b53032c4b1ab8337a0b978dee832ed276dae3815df4e8b5/datasets-4.8.5-py3-none-any.whl", hash = "sha256:5079900781719c0e063a8efdd2cd95a31ad0c63209178669cd23cf1b926149ff", size = 528973, upload-time = "2026-04-27T15:43:53.702Z" }, +] + +[[package]] +name = "dill" +version = "0.4.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/81/e1/56027a71e31b02ddc53c7d65b01e68edf64dea2932122fe7746a516f75d5/dill-0.4.1.tar.gz", hash = "sha256:423092df4182177d4d8ba8290c8a5b640c66ab35ec7da59ccfa00f6fa3eea5fa", size = 187315, upload-time = "2026-01-19T02:36:56.85Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/1e/77/dc8c558f7593132cf8fefec57c4f60c83b16941c574ac5f619abb3ae7933/dill-0.4.1-py3-none-any.whl", hash = "sha256:1e1ce33e978ae97fcfcff5638477032b801c46c7c65cf717f95fbc2248f79a9d", size = 120019, upload-time = "2026-01-19T02:36:55.663Z" }, +] + [[package]] name = "dimos-lerobot-runtime" version = "0.1.0" source = { editable = "." } dependencies = [ - { name = "lerobot" }, + { name = "lerobot", extra = ["dataset", "viz"] }, + { name = "mcap" }, { name = "opencv-contrib-python" }, { name = "transformers", extra = ["torch"] }, ] @@ -240,7 +360,8 @@ tests = [ [package.metadata] requires-dist = [ - { name = "lerobot", specifier = "==0.6.0" }, + { name = "lerobot", extras = ["dataset", "viz"], specifier = "==0.6.0" }, + { name = "mcap", specifier = ">=1.2" }, { name = "opencv-contrib-python", specifier = ">=4.8,<5" }, { name = "transformers", extras = ["torch"], specifier = ">=5.4,<5.6" }, ] @@ -295,6 +416,51 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a7/8e/50f46a9c0ce8d2861a394c1347caae037ea0431d2f67d7feb151cbc4649a/filelock-3.32.3-py3-none-any.whl", hash = "sha256:7f0ca4bcc0e181c60dbbd8aa9ab5b120ebb99e4e064e83636340056f833a1f09", size = 98901, upload-time = "2026-08-13T16:00:03.974Z" }, ] +[[package]] +name = "foxglove-sdk" +version = "0.25.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/ca/25/a0bb15179539bf144f8bf13b176518886191710b40fe4a8173b461d90313/foxglove_sdk-0.25.3.tar.gz", hash = "sha256:32f1066401c37538d478b4a8ecaaae4e9dd019f9f75d3ef47494bfe7d4273dee", size = 549908, upload-time = "2026-06-25T00:24:26.614Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/4f/48/b8d36edbaeef75f52197fd08fe5e431cb73ab32b975377facb468da57b22/foxglove_sdk-0.25.3-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:9667615b342651b3960dd82e6a8403678c1fc701799bf0b28e41be0392eb507d", size = 17812618, upload-time = "2026-06-25T00:24:06.335Z" }, + { url = "https://files.pythonhosted.org/packages/e7/49/2dee8063e2bf2d8cf0db97e732a66a1c75f1faa519ff87f86c9f735d0944/foxglove_sdk-0.25.3-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:d88faf89da6da3d904b6aa1371b809f4d9dcabb8981fcedcfafccdf07b200a49", size = 16379537, upload-time = "2026-06-25T00:24:09.088Z" }, + { url = "https://files.pythonhosted.org/packages/15/3b/2bf924ee1ba5bd9eb49322d8c04850abc444e3323baba425d48112d4cbde/foxglove_sdk-0.25.3-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:dcdcb054c198e892ac271c88e685e292d1ef8f1f21a14e2d3a9de5988474e60a", size = 2261344, upload-time = "2026-06-25T00:24:10.961Z" }, + { url = "https://files.pythonhosted.org/packages/5f/dc/1765ec5b19b0ddecdaff1dad1bc5a332668784e7d99924ebdf003d7e73ad/foxglove_sdk-0.25.3-cp310-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:510c7fe47c45428e9d76eddfae1b9966b7976819a3a194a00f7e8d0f4b3eacc3", size = 2178418, upload-time = "2026-06-25T00:24:12.359Z" }, + { url = "https://files.pythonhosted.org/packages/b2/15/a9182e0d3bbee09261e79249a9d608af4739e6e25791bcd59b23b6522b34/foxglove_sdk-0.25.3-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:18d7dd5e61afc65163438b858a70bec3a30126de1664a97358d4403e20db5eb1", size = 2209079, upload-time = "2026-06-25T00:24:13.575Z" }, + { url = "https://files.pythonhosted.org/packages/df/7b/82cfe8ee2bdef5a0f81c9a9993f587bb322848475f8b94d6d814fef9675a/foxglove_sdk-0.25.3-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:e00cdb87d0a1623cd60a66f7fd531bf492164433e313d4c7f77cb42fdd613420", size = 18620792, upload-time = "2026-06-25T00:24:15.059Z" }, + { url = "https://files.pythonhosted.org/packages/a1/37/6313a16cecd97b0b733967a92f85e4d1929cbbd521433a4db5e070aa850a/foxglove_sdk-0.25.3-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:bcc894b88188d8169973cfbb1370300f671760adea9d6e9447e5a03b2289527d", size = 19220466, upload-time = "2026-06-25T00:24:17.491Z" }, + { url = "https://files.pythonhosted.org/packages/cf/68/120868313bea126b893432f0a0b5143e4042d85946a8aaa4d32039e6484d/foxglove_sdk-0.25.3-cp310-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:139bb1f2ae284673ecb17a1b47955989a58ab3a8639c1fd0f6d3f3afbbbd0fda", size = 2445762, upload-time = "2026-06-25T00:24:19.301Z" }, + { url = "https://files.pythonhosted.org/packages/10/79/3d0744e87c6bd1d72417574df84d94d1b5711c2b0196a3325f7b52c9930d/foxglove_sdk-0.25.3-cp310-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:3d9e671be05e7539a99d84dbd254581616f7ebae8ee8430766788dd6093d58a3", size = 2452345, upload-time = "2026-06-25T00:24:20.511Z" }, + { url = "https://files.pythonhosted.org/packages/53/bd/93dd3cd685b4605877facc12b4f1b7ffb4894cd5040813c9c6fd262e2e55/foxglove_sdk-0.25.3-cp310-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:2dbdef29d2b6115bbbc478f309ac4937fb8765b42681537db99a15282ebaac6f", size = 2468866, upload-time = "2026-06-25T00:24:21.775Z" }, + { url = "https://files.pythonhosted.org/packages/b4/ff/c5f4f8f75c4a6c8b0822da3bf8af8bcb3db27403439dc032c2d7f0b890dc/foxglove_sdk-0.25.3-cp310-abi3-win32.whl", hash = "sha256:a2a1716cade8a9dd842df18f7b1306a9931fa114f485fc552fc14c44ddba8351", size = 1537215, upload-time = "2026-06-25T00:24:22.99Z" }, + { url = "https://files.pythonhosted.org/packages/c7/e1/8ccb4a985c5baf82947e48cff18483c2125ea11e55e8a28950730ab6c065/foxglove_sdk-0.25.3-cp310-abi3-win_amd64.whl", hash = "sha256:ac881ae307ba432766e6141d9098ced2059838226d225fbeb20de1c57d782a9f", size = 16496466, upload-time = "2026-06-25T00:24:24.906Z" }, +] + +[[package]] +name = "frozenlist" +version = "1.8.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/2d/f5/c831fac6cc817d26fd54c7eaccd04ef7e0288806943f7cc5bbf69f3ac1f0/frozenlist-1.8.0.tar.gz", hash = "sha256:3ede829ed8d842f6cd48fc7081d7a41001a56f1f38603f9d49bf3020d59a31ad", size = 45875, upload-time = "2025-10-06T05:38:17.865Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/69/29/948b9aa87e75820a38650af445d2ef2b6b8a6fab1a23b6bb9e4ef0be2d59/frozenlist-1.8.0-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:78f7b9e5d6f2fdb88cdde9440dc147259b62b9d3b019924def9f6478be254ac1", size = 87782, upload-time = "2025-10-06T05:36:06.649Z" }, + { url = "https://files.pythonhosted.org/packages/64/80/4f6e318ee2a7c0750ed724fa33a4bdf1eacdc5a39a7a24e818a773cd91af/frozenlist-1.8.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:229bf37d2e4acdaf808fd3f06e854a4a7a3661e871b10dc1f8f1896a3b05f18b", size = 50594, upload-time = "2025-10-06T05:36:07.69Z" }, + { url = "https://files.pythonhosted.org/packages/2b/94/5c8a2b50a496b11dd519f4a24cb5496cf125681dd99e94c604ccdea9419a/frozenlist-1.8.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:f833670942247a14eafbb675458b4e61c82e002a148f49e68257b79296e865c4", size = 50448, upload-time = "2025-10-06T05:36:08.78Z" }, + { url = "https://files.pythonhosted.org/packages/6a/bd/d91c5e39f490a49df14320f4e8c80161cfcce09f1e2cde1edd16a551abb3/frozenlist-1.8.0-cp312-cp312-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:494a5952b1c597ba44e0e78113a7266e656b9794eec897b19ead706bd7074383", size = 242411, upload-time = "2025-10-06T05:36:09.801Z" }, + { url = "https://files.pythonhosted.org/packages/8f/83/f61505a05109ef3293dfb1ff594d13d64a2324ac3482be2cedc2be818256/frozenlist-1.8.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:96f423a119f4777a4a056b66ce11527366a8bb92f54e541ade21f2374433f6d4", size = 243014, upload-time = "2025-10-06T05:36:11.394Z" }, + { url = "https://files.pythonhosted.org/packages/d8/cb/cb6c7b0f7d4023ddda30cf56b8b17494eb3a79e3fda666bf735f63118b35/frozenlist-1.8.0-cp312-cp312-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:3462dd9475af2025c31cc61be6652dfa25cbfb56cbbf52f4ccfe029f38decaf8", size = 234909, upload-time = "2025-10-06T05:36:12.598Z" }, + { url = "https://files.pythonhosted.org/packages/31/c5/cd7a1f3b8b34af009fb17d4123c5a778b44ae2804e3ad6b86204255f9ec5/frozenlist-1.8.0-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:c4c800524c9cd9bac5166cd6f55285957fcfc907db323e193f2afcd4d9abd69b", size = 250049, upload-time = "2025-10-06T05:36:14.065Z" }, + { url = "https://files.pythonhosted.org/packages/c0/01/2f95d3b416c584a1e7f0e1d6d31998c4a795f7544069ee2e0962a4b60740/frozenlist-1.8.0-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:d6a5df73acd3399d893dafc71663ad22534b5aa4f94e8a2fabfe856c3c1b6a52", size = 256485, upload-time = "2025-10-06T05:36:15.39Z" }, + { url = "https://files.pythonhosted.org/packages/ce/03/024bf7720b3abaebcff6d0793d73c154237b85bdf67b7ed55e5e9596dc9a/frozenlist-1.8.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:405e8fe955c2280ce66428b3ca55e12b3c4e9c336fb2103a4937e891c69a4a29", size = 237619, upload-time = "2025-10-06T05:36:16.558Z" }, + { url = "https://files.pythonhosted.org/packages/69/fa/f8abdfe7d76b731f5d8bd217827cf6764d4f1d9763407e42717b4bed50a0/frozenlist-1.8.0-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:908bd3f6439f2fef9e85031b59fd4f1297af54415fb60e4254a95f75b3cab3f3", size = 250320, upload-time = "2025-10-06T05:36:17.821Z" }, + { url = "https://files.pythonhosted.org/packages/f5/3c/b051329f718b463b22613e269ad72138cc256c540f78a6de89452803a47d/frozenlist-1.8.0-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:294e487f9ec720bd8ffcebc99d575f7eff3568a08a253d1ee1a0378754b74143", size = 246820, upload-time = "2025-10-06T05:36:19.046Z" }, + { url = "https://files.pythonhosted.org/packages/0f/ae/58282e8f98e444b3f4dd42448ff36fa38bef29e40d40f330b22e7108f565/frozenlist-1.8.0-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:74c51543498289c0c43656701be6b077f4b265868fa7f8a8859c197006efb608", size = 250518, upload-time = "2025-10-06T05:36:20.763Z" }, + { url = "https://files.pythonhosted.org/packages/8f/96/007e5944694d66123183845a106547a15944fbbb7154788cbf7272789536/frozenlist-1.8.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:776f352e8329135506a1d6bf16ac3f87bc25b28e765949282dcc627af36123aa", size = 239096, upload-time = "2025-10-06T05:36:22.129Z" }, + { url = "https://files.pythonhosted.org/packages/66/bb/852b9d6db2fa40be96f29c0d1205c306288f0684df8fd26ca1951d461a56/frozenlist-1.8.0-cp312-cp312-win32.whl", hash = "sha256:433403ae80709741ce34038da08511d4a77062aa924baf411ef73d1146e74faf", size = 39985, upload-time = "2025-10-06T05:36:23.661Z" }, + { url = "https://files.pythonhosted.org/packages/b8/af/38e51a553dd66eb064cdf193841f16f077585d4d28394c2fa6235cb41765/frozenlist-1.8.0-cp312-cp312-win_amd64.whl", hash = "sha256:34187385b08f866104f0c0617404c8eb08165ab1272e884abc89c112e9c00746", size = 44591, upload-time = "2025-10-06T05:36:24.958Z" }, + { url = "https://files.pythonhosted.org/packages/a7/06/1dc65480ab147339fecc70797e9c2f69d9cea9cf38934ce08df070fdb9cb/frozenlist-1.8.0-cp312-cp312-win_arm64.whl", hash = "sha256:fe3c58d2f5db5fbd18c2987cba06d51b0529f52bc3a6cdc33d3f4eab725104bd", size = 40102, upload-time = "2025-10-06T05:36:26.333Z" }, + { url = "https://files.pythonhosted.org/packages/9a/9a/e35b4a917281c0b8419d4207f4334c8e8c5dbf4f3f5f9ada73958d937dcc/frozenlist-1.8.0-py3-none-any.whl", hash = "sha256:0c18a16eab41e82c295618a77502e17b195883241c563b00f0aa5106fc4eaa0d", size = 13409, upload-time = "2025-10-06T05:38:16.721Z" }, +] + [[package]] name = "fsspec" version = "2026.2.0" @@ -304,6 +470,11 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e6/ab/fb21f4c939bb440104cc2b396d3be1d9b7a9fd3c6c2a53d98c45b3d7c954/fsspec-2026.2.0-py3-none-any.whl", hash = "sha256:98de475b5cb3bd66bedd5c4679e87b4fdfe1a3bf4d707b151b3c07e58c9a2437", size = 202505, upload-time = "2026-02-05T21:50:51.819Z" }, ] +[package.optional-dependencies] +http = [ + { name = "aiohttp" }, +] + [[package]] name = "gymnasium" version = "1.3.0" @@ -422,6 +593,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/62/a1/3d680cbfd5f4b8f15abc1d571870c5fc3e594bb582bc3b64ea099db13e56/jinja2-3.1.6-py3-none-any.whl", hash = "sha256:85ece4451f492d0c13c5dd7c13a64681a86afae63a5f347908daf103ce6d2f67", size = 134899, upload-time = "2025-03-05T20:05:00.369Z" }, ] +[[package]] +name = "jsonlines" +version = "4.0.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "attrs" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/35/87/bcda8e46c88d0e34cad2f09ee2d0c7f5957bccdb9791b0b934ec84d84be4/jsonlines-4.0.0.tar.gz", hash = "sha256:0c6d2c09117550c089995247f605ae4cf77dd1533041d366351f6f298822ea74", size = 11359, upload-time = "2023-09-01T12:34:44.187Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f8/62/d9ba6323b9202dd2fe166beab8a86d29465c41a0288cbe229fac60c1ab8d/jsonlines-4.0.0-py3-none-any.whl", hash = "sha256:185b334ff2ca5a91362993f42e83588a360cf95ce4b71a73548502bda52a7c55", size = 8701, upload-time = "2023-09-01T12:34:42.563Z" }, +] + [[package]] name = "lerobot" version = "0.6.0" @@ -449,6 +632,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/5d/20/9a96311c19e9d256e65584ca83c49c5782d0f204836e84ceeb420d4d493e/lerobot-0.6.0-py3-none-any.whl", hash = "sha256:b38a564fbc441d98380576863bf68635dde5fc2c42ddc2a39d0486640dc9e9a8", size = 1743768, upload-time = "2026-07-06T10:42:03.165Z" }, ] +[package.optional-dependencies] +dataset = [ + { name = "av" }, + { name = "datasets" }, + { name = "jsonlines" }, + { name = "pandas" }, + { name = "pyarrow" }, + { name = "torchcodec", marker = "sys_platform == 'never'" }, +] +viz = [ + { name = "foxglove-sdk" }, + { name = "rerun-sdk" }, +] + [[package]] name = "librt" version = "0.15.0" @@ -472,6 +669,22 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/5d/ef/dd24f9635c730b86b87587967dda7516b1845e8b17684603d31607fed598/librt-0.15.0-cp312-cp312-win_arm64.whl", hash = "sha256:a54cf9e0ef47b96af580849db5471142200568ce1e02cbf416addab551369570", size = 112292, upload-time = "2026-08-07T10:47:18.222Z" }, ] +[[package]] +name = "lz4" +version = "4.4.5" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/57/51/f1b86d93029f418033dddf9b9f79c8d2641e7454080478ee2aab5123173e/lz4-4.4.5.tar.gz", hash = "sha256:5f0b9e53c1e82e88c10d7c180069363980136b9d7a8306c4dca4f760d60c39f0", size = 172886, upload-time = "2025-11-03T13:02:36.061Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/1b/ac/016e4f6de37d806f7cc8f13add0a46c9a7cfc41a5ddc2bc831d7954cf1ce/lz4-4.4.5-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:df5aa4cead2044bab83e0ebae56e0944cc7fcc1505c7787e9e1057d6d549897e", size = 207163, upload-time = "2025-11-03T13:01:45.895Z" }, + { url = "https://files.pythonhosted.org/packages/8d/df/0fadac6e5bd31b6f34a1a8dbd4db6a7606e70715387c27368586455b7fc9/lz4-4.4.5-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:6d0bf51e7745484d2092b3a51ae6eb58c3bd3ce0300cf2b2c14f76c536d5697a", size = 207150, upload-time = "2025-11-03T13:01:47.205Z" }, + { url = "https://files.pythonhosted.org/packages/b7/17/34e36cc49bb16ca73fb57fbd4c5eaa61760c6b64bce91fcb4e0f4a97f852/lz4-4.4.5-cp312-cp312-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:7b62f94b523c251cf32aa4ab555f14d39bd1a9df385b72443fd76d7c7fb051f5", size = 1292045, upload-time = "2025-11-03T13:01:48.667Z" }, + { url = "https://files.pythonhosted.org/packages/90/1c/b1d8e3741e9fc89ed3b5f7ef5f22586c07ed6bb04e8343c2e98f0fa7ff04/lz4-4.4.5-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:2c3ea562c3af274264444819ae9b14dbbf1ab070aff214a05e97db6896c7597e", size = 1279546, upload-time = "2025-11-03T13:01:50.159Z" }, + { url = "https://files.pythonhosted.org/packages/55/d9/e3867222474f6c1b76e89f3bd914595af69f55bf2c1866e984c548afdc15/lz4-4.4.5-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:24092635f47538b392c4eaeff14c7270d2c8e806bf4be2a6446a378591c5e69e", size = 1368249, upload-time = "2025-11-03T13:01:51.273Z" }, + { url = "https://files.pythonhosted.org/packages/b2/e7/d667d337367686311c38b580d1ca3d5a23a6617e129f26becd4f5dc458df/lz4-4.4.5-cp312-cp312-win32.whl", hash = "sha256:214e37cfe270948ea7eb777229e211c601a3e0875541c1035ab408fbceaddf50", size = 88189, upload-time = "2025-11-03T13:01:52.605Z" }, + { url = "https://files.pythonhosted.org/packages/a5/0b/a54cd7406995ab097fceb907c7eb13a6ddd49e0b231e448f1a81a50af65c/lz4-4.4.5-cp312-cp312-win_amd64.whl", hash = "sha256:713a777de88a73425cf08eb11f742cd2c98628e79a8673d6a52e3c5f0c116f33", size = 99497, upload-time = "2025-11-03T13:01:53.477Z" }, + { url = "https://files.pythonhosted.org/packages/6a/7e/dc28a952e4bfa32ca16fa2eb026e7a6ce5d1411fcd5986cd08c74ec187b9/lz4-4.4.5-cp312-cp312-win_arm64.whl", hash = "sha256:a88cbb729cc333334ccfb52f070463c21560fca63afcf636a9f160a55fac3301", size = 91279, upload-time = "2025-11-03T13:01:54.419Z" }, +] + [[package]] name = "markdown-it-py" version = "4.2.0" @@ -503,6 +716,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e5/f1/216fc1bbfd74011693a4fd837e7026152e89c4bcf3e77b6692fba9923123/markupsafe-3.0.3-cp312-cp312-win_arm64.whl", hash = "sha256:35add3b638a5d900e807944a078b51922212fb3dedb01633a8defc4b01a3c85f", size = 13906, upload-time = "2025-09-27T18:36:40.689Z" }, ] +[[package]] +name = "mcap" +version = "1.4.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "lz4" }, + { name = "zstandard" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/f8/d7/0f17e59733a71bd4d5b38afc8484531c2cd7648b08c79863ee95fc42b002/mcap-1.4.0.tar.gz", hash = "sha256:0528e2f86a61bfec73779e0628e6cf27af83d01d89e20b27d5ec9f0b556a63ac", size = 22155, upload-time = "2026-06-18T21:50:07.965Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/48/14/7e0b2a74b67e16e5f40ab78cc3e5aa4e7bdd55e0aec963d573edc07a15cd/mcap-1.4.0-py3-none-any.whl", hash = "sha256:0b48b1cc951b8d5aabd2599e60d410bae4f1be1819094f54117b7cbf6b3ee2e9", size = 20826, upload-time = "2026-06-18T21:50:06.704Z" }, +] + [[package]] name = "mdurl" version = "0.1.2" @@ -530,6 +756,48 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/43/e3/7d92a15f894aa0c9c4b49b8ee9ac9850d6e63b03c9c32c0367a13ae62209/mpmath-1.3.0-py3-none-any.whl", hash = "sha256:a0b2b9fe80bbcd81a6647ff13108738cfb482d481d826cc0e02f5b35e5c88d2c", size = 536198, upload-time = "2023-03-07T16:47:09.197Z" }, ] +[[package]] +name = "multidict" +version = "6.7.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/1a/c2/c2d94cbe6ac1753f3fc980da97b3d930efe1da3af3c9f5125354436c073d/multidict-6.7.1.tar.gz", hash = "sha256:ec6652a1bee61c53a3e5776b6049172c53b6aaba34f18c9ad04f82712bac623d", size = 102010, upload-time = "2026-01-26T02:46:45.979Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8d/9c/f20e0e2cf80e4b2e4b1c365bf5fe104ee633c751a724246262db8f1a0b13/multidict-6.7.1-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:a90f75c956e32891a4eda3639ce6dd86e87105271f43d43442a3aedf3cddf172", size = 76893, upload-time = "2026-01-26T02:43:52.754Z" }, + { url = "https://files.pythonhosted.org/packages/fe/cf/18ef143a81610136d3da8193da9d80bfe1cb548a1e2d1c775f26b23d024a/multidict-6.7.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:3fccb473e87eaa1382689053e4a4618e7ba7b9b9b8d6adf2027ee474597128cd", size = 45456, upload-time = "2026-01-26T02:43:53.893Z" }, + { url = "https://files.pythonhosted.org/packages/a9/65/1caac9d4cd32e8433908683446eebc953e82d22b03d10d41a5f0fefe991b/multidict-6.7.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:b0fa96985700739c4c7853a43c0b3e169360d6855780021bfc6d0f1ce7c123e7", size = 43872, upload-time = "2026-01-26T02:43:55.041Z" }, + { url = "https://files.pythonhosted.org/packages/cf/3b/d6bd75dc4f3ff7c73766e04e705b00ed6dbbaccf670d9e05a12b006f5a21/multidict-6.7.1-cp312-cp312-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:cb2a55f408c3043e42b40cc8eecd575afa27b7e0b956dfb190de0f8499a57a53", size = 251018, upload-time = "2026-01-26T02:43:56.198Z" }, + { url = "https://files.pythonhosted.org/packages/fd/80/c959c5933adedb9ac15152e4067c702a808ea183a8b64cf8f31af8ad3155/multidict-6.7.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:eb0ce7b2a32d09892b3dd6cc44877a0d02a33241fafca5f25c8b6b62374f8b75", size = 258883, upload-time = "2026-01-26T02:43:57.499Z" }, + { url = "https://files.pythonhosted.org/packages/86/85/7ed40adafea3d4f1c8b916e3b5cc3a8e07dfcdcb9cd72800f4ed3ca1b387/multidict-6.7.1-cp312-cp312-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:c3a32d23520ee37bf327d1e1a656fec76a2edd5c038bf43eddfa0572ec49c60b", size = 242413, upload-time = "2026-01-26T02:43:58.755Z" }, + { url = "https://files.pythonhosted.org/packages/d2/57/b8565ff533e48595503c785f8361ff9a4fde4d67de25c207cd0ba3befd03/multidict-6.7.1-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:9c90fed18bffc0189ba814749fdcc102b536e83a9f738a9003e569acd540a733", size = 268404, upload-time = "2026-01-26T02:44:00.216Z" }, + { url = "https://files.pythonhosted.org/packages/e0/50/9810c5c29350f7258180dfdcb2e52783a0632862eb334c4896ac717cebcb/multidict-6.7.1-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:da62917e6076f512daccfbbde27f46fed1c98fee202f0559adec8ee0de67f71a", size = 269456, upload-time = "2026-01-26T02:44:02.202Z" }, + { url = "https://files.pythonhosted.org/packages/f3/8d/5e5be3ced1d12966fefb5c4ea3b2a5b480afcea36406559442c6e31d4a48/multidict-6.7.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:bfde23ef6ed9db7eaee6c37dcec08524cb43903c60b285b172b6c094711b3961", size = 256322, upload-time = "2026-01-26T02:44:03.56Z" }, + { url = "https://files.pythonhosted.org/packages/31/6e/d8a26d81ac166a5592782d208dd90dfdc0a7a218adaa52b45a672b46c122/multidict-6.7.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:3758692429e4e32f1ba0df23219cd0b4fc0a52f476726fff9337d1a57676a582", size = 253955, upload-time = "2026-01-26T02:44:04.845Z" }, + { url = "https://files.pythonhosted.org/packages/59/4c/7c672c8aad41534ba619bcd4ade7a0dc87ed6b8b5c06149b85d3dd03f0cd/multidict-6.7.1-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:398c1478926eca669f2fd6a5856b6de9c0acf23a2cb59a14c0ba5844fa38077e", size = 251254, upload-time = "2026-01-26T02:44:06.133Z" }, + { url = "https://files.pythonhosted.org/packages/7b/bd/84c24de512cbafbdbc39439f74e967f19570ce7924e3007174a29c348916/multidict-6.7.1-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:c102791b1c4f3ab36ce4101154549105a53dc828f016356b3e3bcae2e3a039d3", size = 252059, upload-time = "2026-01-26T02:44:07.518Z" }, + { url = "https://files.pythonhosted.org/packages/fa/ba/f5449385510825b73d01c2d4087bf6d2fccc20a2d42ac34df93191d3dd03/multidict-6.7.1-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:a088b62bd733e2ad12c50dad01b7d0166c30287c166e137433d3b410add807a6", size = 263588, upload-time = "2026-01-26T02:44:09.382Z" }, + { url = "https://files.pythonhosted.org/packages/d7/11/afc7c677f68f75c84a69fe37184f0f82fce13ce4b92f49f3db280b7e92b3/multidict-6.7.1-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:3d51ff4785d58d3f6c91bdbffcb5e1f7ddfda557727043aa20d20ec4f65e324a", size = 259642, upload-time = "2026-01-26T02:44:10.73Z" }, + { url = "https://files.pythonhosted.org/packages/2b/17/ebb9644da78c4ab36403739e0e6e0e30ebb135b9caf3440825001a0bddcb/multidict-6.7.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:fc5907494fccf3e7d3f94f95c91d6336b092b5fc83811720fae5e2765890dfba", size = 251377, upload-time = "2026-01-26T02:44:12.042Z" }, + { url = "https://files.pythonhosted.org/packages/ca/a4/840f5b97339e27846c46307f2530a2805d9d537d8b8bd416af031cad7fa0/multidict-6.7.1-cp312-cp312-win32.whl", hash = "sha256:28ca5ce2fd9716631133d0e9a9b9a745ad7f60bac2bccafb56aa380fc0b6c511", size = 41887, upload-time = "2026-01-26T02:44:14.245Z" }, + { url = "https://files.pythonhosted.org/packages/80/31/0b2517913687895f5904325c2069d6a3b78f66cc641a86a2baf75a05dcbb/multidict-6.7.1-cp312-cp312-win_amd64.whl", hash = "sha256:fcee94dfbd638784645b066074b338bc9cc155d4b4bffa4adce1615c5a426c19", size = 46053, upload-time = "2026-01-26T02:44:15.371Z" }, + { url = "https://files.pythonhosted.org/packages/0c/5b/aba28e4ee4006ae4c7df8d327d31025d760ffa992ea23812a601d226e682/multidict-6.7.1-cp312-cp312-win_arm64.whl", hash = "sha256:ba0a9fb644d0c1a2194cf7ffb043bd852cea63a57f66fbd33959f7dae18517bf", size = 43307, upload-time = "2026-01-26T02:44:16.852Z" }, + { url = "https://files.pythonhosted.org/packages/81/08/7036c080d7117f28a4af526d794aab6a84463126db031b007717c1a6676e/multidict-6.7.1-py3-none-any.whl", hash = "sha256:55d97cc6dae627efa6a6e548885712d4864b81110ac76fa4e534c03819fa4a56", size = 12319, upload-time = "2026-01-26T02:46:44.004Z" }, +] + +[[package]] +name = "multiprocess" +version = "0.70.19" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "dill" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a2/f2/e783ac7f2aeeed14e9e12801f22529cc7e6b7ab80928d6dcce4e9f00922d/multiprocess-0.70.19.tar.gz", hash = "sha256:952021e0e6c55a4a9fe4cd787895b86e239a40e76802a789d6305398d3975897", size = 2079989, upload-time = "2026-01-19T06:47:39.744Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e3/45/8004d1e6b9185c1a444d6b55ac5682acf9d98035e54386d967366035a03a/multiprocess-0.70.19-py310-none-any.whl", hash = "sha256:97404393419dcb2a8385910864eedf47a3cadf82c66345b44f036420eb0b5d87", size = 134948, upload-time = "2026-01-19T06:47:32.325Z" }, + { url = "https://files.pythonhosted.org/packages/86/c2/dec9722dc3474c164a0b6bcd9a7ed7da542c98af8cabce05374abab35edd/multiprocess-0.70.19-py311-none-any.whl", hash = "sha256:928851ae7973aea4ce0eaf330bbdafb2e01398a91518d5c8818802845564f45c", size = 144457, upload-time = "2026-01-19T06:47:33.711Z" }, + { url = "https://files.pythonhosted.org/packages/71/70/38998b950a97ea279e6bd657575d22d1a2047256caf707d9a10fbce4f065/multiprocess-0.70.19-py312-none-any.whl", hash = "sha256:3a56c0e85dd5025161bac5ce138dcac1e49174c7d8e74596537e729fd5c53c28", size = 150281, upload-time = "2026-01-19T06:47:35.037Z" }, + { url = "https://files.pythonhosted.org/packages/7e/82/69e539c4c2027f1e1697e09aaa2449243085a0edf81ae2c6341e84d769b6/multiprocess-0.70.19-py39-none-any.whl", hash = "sha256:0d4b4397ed669d371c81dcd1ef33fd384a44d6c3de1bd0ca7ac06d837720d3c5", size = 133477, upload-time = "2026-01-19T06:47:38.619Z" }, +] + [[package]] name = "mypy" version = "1.19.0" @@ -773,6 +1041,27 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/20/12/38679034af332785aac8774540895e234f4d07f7545804097de4b666afd8/packaging-25.0-py3-none-any.whl", hash = "sha256:29572ef2b1f17581046b3a2227d5c611fb25ec70ca1ba8554b24b0e69331a484", size = 66469, upload-time = "2025-04-19T11:48:57.875Z" }, ] +[[package]] +name = "pandas" +version = "2.3.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, + { name = "python-dateutil" }, + { name = "pytz" }, + { name = "tzdata" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/33/01/d40b85317f86cf08d853a4f495195c73815fdf205eef3993821720274518/pandas-2.3.3.tar.gz", hash = "sha256:e05e1af93b977f7eafa636d043f9f94c7ee3ac81af99c13508215942e64c993b", size = 4495223, upload-time = "2025-09-29T23:34:51.853Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9c/fb/231d89e8637c808b997d172b18e9d4a4bc7bf31296196c260526055d1ea0/pandas-2.3.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:6d21f6d74eb1725c2efaa71a2bfc661a0689579b58e9c0ca58a739ff0b002b53", size = 11597846, upload-time = "2025-09-29T23:19:48.856Z" }, + { url = "https://files.pythonhosted.org/packages/5c/bd/bf8064d9cfa214294356c2d6702b716d3cf3bb24be59287a6a21e24cae6b/pandas-2.3.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:3fd2f887589c7aa868e02632612ba39acb0b8948faf5cc58f0850e165bd46f35", size = 10729618, upload-time = "2025-09-29T23:39:08.659Z" }, + { url = "https://files.pythonhosted.org/packages/57/56/cf2dbe1a3f5271370669475ead12ce77c61726ffd19a35546e31aa8edf4e/pandas-2.3.3-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ecaf1e12bdc03c86ad4a7ea848d66c685cb6851d807a26aa245ca3d2017a1908", size = 11737212, upload-time = "2025-09-29T23:19:59.765Z" }, + { url = "https://files.pythonhosted.org/packages/e5/63/cd7d615331b328e287d8233ba9fdf191a9c2d11b6af0c7a59cfcec23de68/pandas-2.3.3-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b3d11d2fda7eb164ef27ffc14b4fcab16a80e1ce67e9f57e19ec0afaf715ba89", size = 12362693, upload-time = "2025-09-29T23:20:14.098Z" }, + { url = "https://files.pythonhosted.org/packages/a6/de/8b1895b107277d52f2b42d3a6806e69cfef0d5cf1d0ba343470b9d8e0a04/pandas-2.3.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:a68e15f780eddf2b07d242e17a04aa187a7ee12b40b930bfdd78070556550e98", size = 12771002, upload-time = "2025-09-29T23:20:26.76Z" }, + { url = "https://files.pythonhosted.org/packages/87/21/84072af3187a677c5893b170ba2c8fbe450a6ff911234916da889b698220/pandas-2.3.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:371a4ab48e950033bcf52b6527eccb564f52dc826c02afd9a1bc0ab731bba084", size = 13450971, upload-time = "2025-09-29T23:20:41.344Z" }, + { url = "https://files.pythonhosted.org/packages/86/41/585a168330ff063014880a80d744219dbf1dd7a1c706e75ab3425a987384/pandas-2.3.3-cp312-cp312-win_amd64.whl", hash = "sha256:a16dcec078a01eeef8ee61bf64074b4e524a2a3f4b3be9326420cabe59c4778b", size = 10992722, upload-time = "2025-09-29T23:20:54.139Z" }, +] + [[package]] name = "pathspec" version = "1.1.1" @@ -808,6 +1097,32 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" }, ] +[[package]] +name = "propcache" +version = "0.5.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/ec/44/c87281c333769159c50594f22610f77398a47ccbfbbf23074e744e86f87c/propcache-0.5.2.tar.gz", hash = "sha256:01c4fc7480cd0598bb4b57022df55b9ca296da7fc5a8760bd8451a7e63a7d427", size = 50208, upload-time = "2026-05-08T21:02:12.199Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/4a/cb/e27bc2b2737a0bb49962b275efa051e8f1c35a936df7d5139b6b658b7dc9/propcache-0.5.2-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:806719138ecd720339a12410fb9614ac9b2b2d3a5fdf8235d56981c36f4039ba", size = 95887, upload-time = "2026-05-08T21:00:11.277Z" }, + { url = "https://files.pythonhosted.org/packages/e6/13/b8ae04c59392f8d11c6cd9fb4011d1dc7c86b81225c770280300e259ffe1/propcache-0.5.2-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:db2b80ea58eab4f86b2beec3cc8b39e8ff9276ac20e96b7cce43c8ae84cd6b5a", size = 54654, upload-time = "2026-05-08T21:00:12.604Z" }, + { url = "https://files.pythonhosted.org/packages/2c/7d/49777a3e20b55863d4794384a38acd460c04157b0a00f8602b0d508b8431/propcache-0.5.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:e5cbfac9f61484f7e9f3597775500cd3ebe8274e9b050c38f9525c77c97520bf", size = 55190, upload-time = "2026-05-08T21:00:13.935Z" }, + { url = "https://files.pythonhosted.org/packages/44/c7/085d0cd63062e84044e3f05797749c3f8e3938ff3aeb0eb2f69d43fafc91/propcache-0.5.2-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5dbc581d2814337da56222fab8dc5f161cd798a434e49bac27930aaef798e144", size = 59995, upload-time = "2026-05-08T21:00:15.526Z" }, + { url = "https://files.pythonhosted.org/packages/9c/42/32cf8e3009e92b2645cf1e944f701e8ea4e924dffde1ee26db860bcbf7e4/propcache-0.5.2-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:857187f381f88c8e2fa2fe56ab94879d011b883d5a2ee5a1b60a8cd2a06846d9", size = 63422, upload-time = "2026-05-08T21:00:16.824Z" }, + { url = "https://files.pythonhosted.org/packages/9e/1b/f112433f99fc979431b87a39ef169e3f8df070d99a72792c56d6937ac48b/propcache-0.5.2-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:178b4a2cdaac1818e2bf1c5a99b94383fa73ea5382e032a48dec07dc5668dc42", size = 64342, upload-time = "2026-05-08T21:00:18.362Z" }, + { url = "https://files.pythonhosted.org/packages/14/15/5574111ae50dd6e879456888c0eadd4c5a869959775854e18e18a6b345f3/propcache-0.5.2-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6f328175a2cde1f0ff2c4ed8ce968b9dcfb55f3a7153f39e2957ed994da13476", size = 61639, upload-time = "2026-05-08T21:00:19.692Z" }, + { url = "https://files.pythonhosted.org/packages/cc/da/4d775080b1490c0ae604acda868bd71aabe3a89ed16f2aa4339eb8a283e7/propcache-0.5.2-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:5671d09a36b06d0fd4a3da0fccbcae360e9b1570924171a15e9e0997f0249fba", size = 61588, upload-time = "2026-05-08T21:00:21.155Z" }, + { url = "https://files.pythonhosted.org/packages/04/ac/f076982cbe2195ee9cf32de5a1e46951d9fb399fc207f390562dd0fd8fb2/propcache-0.5.2-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:80168e2ebe4d3ec6599d10ad8f520304ae1cad9b6c5a95372aef1b66b7bfb53a", size = 60029, upload-time = "2026-05-08T21:00:22.713Z" }, + { url = "https://files.pythonhosted.org/packages/70/60/189be62e0dd898dce3b331e1b8c7a543cd3a405ac0c81fe8ee8a9d5d77e1/propcache-0.5.2-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:45f11346f884bc47444f6e6647131055844134c3175b629f84952e2b5cd62b64", size = 56774, upload-time = "2026-05-08T21:00:24.001Z" }, + { url = "https://files.pythonhosted.org/packages/ea/9e/93377b9c7939c1ffae98f878dee955efadfd638078bc86dbc21f9d52f651/propcache-0.5.2-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:8e778ebd44ef4f66ed60a0416b06b489687db264a9c0b3620362f26489492913", size = 63532, upload-time = "2026-05-08T21:00:25.545Z" }, + { url = "https://files.pythonhosted.org/packages/14/f9/590ef6cfb9b8028d516d287812ece32bb0bc5f11fbb9c8bf6b2e6313fec8/propcache-0.5.2-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:c0cb9ed24c8964e172768d455a38254c2dd8a552905729ce006cad3d3dda59b1", size = 61592, upload-time = "2026-05-08T21:00:27.186Z" }, + { url = "https://files.pythonhosted.org/packages/b4/5e/70958b3034c297a630bba2f17ca7abc2d5f39a803ad7e370ab79d1ecd022/propcache-0.5.2-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:1d1ad32d9d4355e2be65574fd0bfd3677e7066b009cd5b9b2dee8aa6a6393b33", size = 64788, upload-time = "2026-05-08T21:00:28.8Z" }, + { url = "https://files.pythonhosted.org/packages/12/fd/77fe5936d8c3086ca9048f7f415f122ed82e53884a9ec193646b42deef06/propcache-0.5.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:c80f4ba3e8f00189165999a742ee526ebeccedf6c3f7beb0c7df821e9772435a", size = 62514, upload-time = "2026-05-08T21:00:30.098Z" }, + { url = "https://files.pythonhosted.org/packages/cf/74/66bd798b5b3be70aa1b391f5cc9d6a0a5532d7fd3b19ec0b213e72e6ad9d/propcache-0.5.2-cp312-cp312-win32.whl", hash = "sha256:8c7972d8f193740d9175f0998ab38717e6cd322d5935c5b0fef8c0d323fd9031", size = 39018, upload-time = "2026-05-08T21:00:31.622Z" }, + { url = "https://files.pythonhosted.org/packages/61/7c/5c0d34aa3024694d6dcb9271cdbdd08c4e47c1c0ad95ec7e7bc74cdea145/propcache-0.5.2-cp312-cp312-win_amd64.whl", hash = "sha256:d9ee8826a7d47863a08ac44e1a5f611a462eefc3a194b492da242128bec75b42", size = 42322, upload-time = "2026-05-08T21:00:32.918Z" }, + { url = "https://files.pythonhosted.org/packages/4d/91/875812f1a3feb20ceba818ef39fbe4d92f1081e04ac815c822496d0d038b/propcache-0.5.2-cp312-cp312-win_arm64.whl", hash = "sha256:2800a4a8ead6b28cccd1ec54b59346f0def7922ee1c7598e8499c733cfbb7c84", size = 38172, upload-time = "2026-05-08T21:00:35.124Z" }, + { url = "https://files.pythonhosted.org/packages/3a/ed/1cdcab6ba3d6ab7feca11fc14f0eeea80755bb53ef4e892079f31b10a25f/propcache-0.5.2-py3-none-any.whl", hash = "sha256:be1ddfcbb376e3de5d2e2db1d58d6d67463e6b4f9f040c000de8e300295465fe", size = 14036, upload-time = "2026-05-08T21:02:10.673Z" }, +] + [[package]] name = "psutil" version = "7.2.2" @@ -824,6 +1139,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/8c/c7/7bb2e321574b10df20cbde462a94e2b71d05f9bbda251ef27d104668306a/psutil-7.2.2-cp37-abi3-win_arm64.whl", hash = "sha256:8c233660f575a5a89e6d4cb65d9f938126312bca76d8fe087b947b3a1aaac9ee", size = 134617, upload-time = "2026-01-28T18:15:36.514Z" }, ] +[[package]] +name = "pyarrow" +version = "25.0.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/3d/e3/27f57f80141379d60defe6703eb50a707325706f07fedfd1312c7a751995/pyarrow-25.0.1.tar.gz", hash = "sha256:9150a83248bfed9813ea3c3af74c3856c1984d444aa28e58bf7733b9750ddf6a", size = 1201653, upload-time = "2026-08-10T12:40:53.904Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a6/e2/9ab15b88cbfac28e16419ce5439ec29234c5172cb8259301b4ba639bdec0/pyarrow-25.0.1-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:df961f2e7ae9cf496459259d798652c70625f6c080650d6952f8c04053c58ee9", size = 35861559, upload-time = "2026-08-10T12:38:02.567Z" }, + { url = "https://files.pythonhosted.org/packages/58/79/a0036dbe1eabe1f73127427342f1d99982584c4a2cde2651d6c93499c6f6/pyarrow-25.0.1-cp312-cp312-macosx_12_0_x86_64.whl", hash = "sha256:cc4aa407fde9fc660be3939e49ea31f50f3e9fec17c0ec63159f7711edd3efc9", size = 37628383, upload-time = "2026-08-10T12:38:09.083Z" }, + { url = "https://files.pythonhosted.org/packages/13/49/d93a57d375f4bf0cf82913dd6bb54acafde83dd993be2282c81ac5616cad/pyarrow-25.0.1-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:4340f0ba6c1d2e13f21658de1d7c662ca2545018568d0030a1e9afca159d87e3", size = 46820190, upload-time = "2026-08-10T12:38:15.458Z" }, + { url = "https://files.pythonhosted.org/packages/60/c9/711ca85d79f1ec98f29a5eae2b051e25b4ecec5de3e3c0e2d5c5dcb15664/pyarrow-25.0.1-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:5389cdf79447ed1515c9e31620e6e1e2302249564d603f2ad727d4f6d313e4c3", size = 50102437, upload-time = "2026-08-10T12:38:22.487Z" }, + { url = "https://files.pythonhosted.org/packages/80/53/8fb8359ff17cfb6263a1cf3ebf7caec9fe197de118719e84fcb1d0618026/pyarrow-25.0.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:d51592cb7561e87877c506113e7adbf1342ab579e6c21f0ef44b8ba41cb74c80", size = 49942424, upload-time = "2026-08-10T12:38:28.755Z" }, + { url = "https://files.pythonhosted.org/packages/e8/83/4e5ae02a9341571b18a6fca380ac7a58ce6ddae7ab3c060208c0a1e79f02/pyarrow-25.0.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:6109c94d8b9f3b17a041daca16cacb2f651ad8f1ef70a4232c2c0f37a23da2a8", size = 53144206, upload-time = "2026-08-10T12:38:34.862Z" }, + { url = "https://files.pythonhosted.org/packages/65/ee/197cbf47e49f83e6ebeb946a5259a48a638dea27ac774db42fe78022179d/pyarrow-25.0.1-cp312-cp312-win_amd64.whl", hash = "sha256:8858d7bfc22e3f51529aeaa4077225029724623e4595dc9eff8c793935c34140", size = 27953934, upload-time = "2026-08-10T12:38:39.808Z" }, +] + [[package]] name = "pygments" version = "2.20.0" @@ -860,6 +1190,27 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/5a/cc/06253936f4a7fa2e0f48dfe6d851d9c56df896a9ab09ac019d70b760619c/pytest_mock-3.15.1-py3-none-any.whl", hash = "sha256:0a25e2eb88fe5168d535041d09a4529a188176ae608a6d249ee65abc0949630d", size = 10095, upload-time = "2025-09-16T16:37:25.734Z" }, ] +[[package]] +name = "python-dateutil" +version = "2.9.0.post0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "six" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/66/c0/0c8b6ad9f17a802ee498c46e004a0eb49bc148f2fd230864601a86dcf6db/python-dateutil-2.9.0.post0.tar.gz", hash = "sha256:37dd54208da7e1cd875388217d5e00ebd4179249f90fb72437e91a35459a0ad3", size = 342432, upload-time = "2024-03-01T18:36:20.211Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl", hash = "sha256:a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427", size = 229892, upload-time = "2024-03-01T18:36:18.57Z" }, +] + +[[package]] +name = "pytz" +version = "2026.3.post1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/fb/48/fb042503b6ca6cd271261dc559fd6432f7d8c713153e9ec5c591af4dfc1c/pytz-2026.3.post1.tar.gz", hash = "sha256:2211d3fcf9a797d3405cac96ac7f61d80e6a644f72a3309607282fe8a2010c5d", size = 319745, upload-time = "2026-07-25T15:12:07.385Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0f/7b/39c34ca613b0b198cb866466651b26b045e2009864c5183c979a3b83f383/pytz-2026.3.post1-py2.py3-none-any.whl", hash = "sha256:dd95840dd199baea12d9cc096a1d452caa6596a1c1e4b5f3dbd1541855d5e815", size = 508283, upload-time = "2026-07-25T15:12:05.782Z" }, +] + [[package]] name = "pyyaml" version = "6.0.3" @@ -929,6 +1280,25 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a0/f4/c67b0b3f1b9245e8d266f0f112c500d50e5b4e83cb6f3b71b6528104182a/requests-2.34.2-py3-none-any.whl", hash = "sha256:2a0d60c172f83ac6ab31e4554906c0f3b3588d37b5cb939b1c061f4907e278e0", size = 73075, upload-time = "2026-05-14T19:25:26.443Z" }, ] +[[package]] +name = "rerun-sdk" +version = "0.33.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "attrs" }, + { name = "numpy" }, + { name = "pillow" }, + { name = "psutil" }, + { name = "pyarrow" }, + { name = "typing-extensions" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/16/07/380198590b194f1c17052d672865aa4a56e606eae47665f66edfb391999d/rerun_sdk-0.33.1-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:c0115b710289022587bd2e9ecb715f98d1e87dd7d8ac48e053324131d4addc89", size = 125706253, upload-time = "2026-06-22T09:04:18.974Z" }, + { url = "https://files.pythonhosted.org/packages/f8/eb/6741bbf6868175ab126aff58d372066241c6cd2fc1c4f82ed64069728e73/rerun_sdk-0.33.1-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:14451d31bc7bd0f7c6bcc9c1213ed679ab81b65ecd8b36eae99f738219897dc5", size = 135278374, upload-time = "2026-06-22T09:04:26.768Z" }, + { url = "https://files.pythonhosted.org/packages/04/28/fd3b832652900fe3415739e7411c8af8af4c44e9e1a9d55d79e37f7f9094/rerun_sdk-0.33.1-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:0f87da7a270614074aca37846350f9e257f65081345474748f578c7da64fdeba", size = 139565470, upload-time = "2026-06-22T09:04:34.941Z" }, + { url = "https://files.pythonhosted.org/packages/19/41/339920f5a6734054c07bcae543365a7ef3368ceee3eb67906e2e38bd1d67/rerun_sdk-0.33.1-cp310-abi3-win_amd64.whl", hash = "sha256:b2c2af67f3c2a85b282669d97b52596593fcfdd19bf57c423f18827c837a6e49", size = 120411717, upload-time = "2026-06-22T09:04:42.643Z" }, +] + [[package]] name = "rich" version = "15.0.0" @@ -984,6 +1354,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e0/f9/0595336914c5619e5f28a1fb793285925a8cd4b432c9da0a987836c7f822/shellingham-1.5.4-py2.py3-none-any.whl", hash = "sha256:7ecfff8f2fd72616f7481040475a65b2bf8af90a56c89140852d1120324e8686", size = 9755, upload-time = "2023-10-24T04:13:38.866Z" }, ] +[[package]] +name = "six" +version = "1.17.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/94/e7/b2c673351809dca68a0e064b6af791aa332cf192da575fd474ed7d6f16a2/six-1.17.0.tar.gz", hash = "sha256:ff70335d468e7eb6ec65b95b99d3a2836546063f63acc5171de367e834932a81", size = 34031, upload-time = "2024-12-04T17:35:28.174Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274", size = 11050, upload-time = "2024-12-04T17:35:26.475Z" }, +] + [[package]] name = "sympy" version = "1.14.0" @@ -1067,6 +1446,11 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/1c/ff/6756f1c7ee302f6d202120e0f4f05b432b839908f9071157302cedfc5232/torch-2.11.0-cp312-cp312-win_amd64.whl", hash = "sha256:fbf39280699d1b869f55eac536deceaa1b60bd6788ba74f399cc67e60a5fab10", size = 114556047, upload-time = "2026-03-23T18:10:55.931Z" }, ] +[[package]] +name = "torchcodec" +version = "0.11.1" +source = { registry = "https://pypi.org/simple" } + [[package]] name = "torchvision" version = "0.26.0" @@ -1167,6 +1551,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/65/f3/107a22063bf27bdccf2024833d3445f4eea42b2e598abfbd46f6a63b6cb0/typing_inspect-0.9.0-py3-none-any.whl", hash = "sha256:9ee6fc59062311ef8547596ab6b955e1b8aa46242d854bfc78f4f6b0eff35f9f", size = 8827, upload-time = "2023-05-24T20:25:45.287Z" }, ] +[[package]] +name = "tzdata" +version = "2026.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/92/ff/5a28bdfd8c3ebec42564ac7d0e54ca3db65044a9314a97f9564fa7a1e926/tzdata-2026.3.tar.gz", hash = "sha256:4a1518b8993086a7982523e071643f3c0e5f213e75b21318e78bcabfff9d1415", size = 198674, upload-time = "2026-07-10T08:50:37.887Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e5/6d/b53b99a9f2766d095985947a5782f1702cabb129a34f7a802d7197af832f/tzdata-2026.3-py2.py3-none-any.whl", hash = "sha256:dc096730c87af6cab1b171c9d532be840741ff5d459015e7f6947bd7d7e54931", size = 348168, upload-time = "2026-07-10T08:50:36.46Z" }, +] + [[package]] name = "urllib3" version = "2.7.0" @@ -1175,3 +1568,92 @@ sdist = { url = "https://files.pythonhosted.org/packages/53/0c/06f8b233b8fd13b9e wheels = [ { url = "https://files.pythonhosted.org/packages/7f/3e/5db95bcf282c52709639744ca2a8b149baccf648e39c8cc87553df9eae0c/urllib3-2.7.0-py3-none-any.whl", hash = "sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897", size = 131087, upload-time = "2026-05-07T16:13:17.151Z" }, ] + +[[package]] +name = "xxhash" +version = "4.0.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f6/a5/1386f35da1475fcaeef42581deae73417c6d2a6a0b2d2e8914de18844dcd/xxhash-4.0.1.tar.gz", hash = "sha256:d55bf4ef10eb09b8b6866790e083d26d087d84caa3cc0946ba87c3ca7ecaf7b7", size = 101513, upload-time = "2026-08-17T08:24:08.557Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/26/6c/dc7cffeadd06336cd934947187cd38abb263103bbc552ca0f55fe4ff595a/xxhash-4.0.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:1ee523f51718e41753f04f7102bb4dc55a18d2ea5cbaceef8ec7ca08571bd428", size = 38444, upload-time = "2026-08-17T08:21:54.332Z" }, + { url = "https://files.pythonhosted.org/packages/75/c9/cf736f6db8c3273af18925061572db0d4357818a9ce425f4b5fb0021918e/xxhash-4.0.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:515a822c73abbf6a0b7c70976d9662be342835c9d78b8dc7c023411f39c35dbc", size = 36195, upload-time = "2026-08-17T08:35:13.004Z" }, + { url = "https://files.pythonhosted.org/packages/da/a2/ca1929354b6851529d0148f7f335b5e2b0281f83bab3e19f0896dc579796/xxhash-4.0.1-cp312-cp312-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:f5d031f35962e5483a613214e61f09fe24ab523062c3646d592dc16c4a217451", size = 253113, upload-time = "2026-08-17T08:20:52.152Z" }, + { url = "https://files.pythonhosted.org/packages/de/bb/542005206af59518bc8d78a210f1e0172217bc53beb32f64a5b632e72b6b/xxhash-4.0.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:da0264844a09b538c894e5eff25313d941deb4dedec2131b98418a71a3c9944e", size = 276525, upload-time = "2026-08-17T08:21:01.886Z" }, + { url = "https://files.pythonhosted.org/packages/1b/df/607cff25dcb0f1d35c3b04493f6ad8471edb03fd4eacbdcc5ceddef1f3e9/xxhash-4.0.1-cp312-cp312-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:1642907941ee4b75aacc3db688af52ea02ca2305ab22af7ee686ed726b332684", size = 297703, upload-time = "2026-08-17T08:21:57.958Z" }, + { url = "https://files.pythonhosted.org/packages/15/ba/9d2275eea0b9d9c6b02921be23f7588356c60df95c763b25f0e045894d43/xxhash-4.0.1-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:4af350bc3f329970c0e3a59af84a8a30998bf8a9167eb50cd48e59baaa1d7bec", size = 280252, upload-time = "2026-08-17T08:20:47.299Z" }, + { url = "https://files.pythonhosted.org/packages/1d/aa/2299d9f6369e550aef2abb64945e39daa34412725aa46a20d99b74d76f67/xxhash-4.0.1-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:8ba782ca3bf1e81492611152b9a0d5264971339e95e34d69de0ac2c926be496d", size = 511041, upload-time = "2026-08-17T08:20:36.771Z" }, + { url = "https://files.pythonhosted.org/packages/83/97/31bd8b8279e6935a0719f6910ced15e9d5a2cd554b253f6027ce1b5a1c2c/xxhash-4.0.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:237b8f63a2a0fcfb1ffc06e21dad23add44e6d354b2b014364a1d41e419a4dee", size = 261812, upload-time = "2026-08-17T08:22:00.469Z" }, + { url = "https://files.pythonhosted.org/packages/2d/c1/d180a2da23c105d8e0b02d54f9f5841013fc81c233010ec781e31f1aee4c/xxhash-4.0.1-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:81507a68ba84c55241fb61cce1469f473a5da4205fc8ef6f698e5948eea8dd88", size = 339878, upload-time = "2026-08-17T08:35:17.626Z" }, + { url = "https://files.pythonhosted.org/packages/a8/3d/f584cd3172fe934f0f5a0a3917d0d7ce781f74d794fd43bb72be71c3ef6f/xxhash-4.0.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:5f1ea31d61bcd2cd2f3ec4ca80a64187bbd7948f490b63cf0dcbc6e717b4c1e9", size = 272871, upload-time = "2026-08-17T08:20:56.067Z" }, + { url = "https://files.pythonhosted.org/packages/34/50/2c7956b2b551682e00b9aebce9ceb0a991a131d65f9850c09f5f9760be2e/xxhash-4.0.1-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:06713a5aaf1d0905c5579416c020c02e42b3ceb931e86c7d3b7fb85403dee3f3", size = 301440, upload-time = "2026-08-17T08:21:35.911Z" }, + { url = "https://files.pythonhosted.org/packages/eb/a2/0739f6482184a8026f4b022718f5f815d352059312e80696825433f0a8e7/xxhash-4.0.1-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:e8cda075b10bb3917b002c74a04f9e02b7d13b5bf732571404d51c52b11c7329", size = 260157, upload-time = "2026-08-17T08:22:01.416Z" }, + { url = "https://files.pythonhosted.org/packages/a1/25/b31a7bcf1d7d116842812e54f9b944843b4236ea4fa85634e8259f342212/xxhash-4.0.1-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:c10b9206753b64aa791b35b201485477525b26fdec5bf86e8364c388a03e2592", size = 278233, upload-time = "2026-08-17T08:21:15.674Z" }, + { url = "https://files.pythonhosted.org/packages/db/e8/5293bae090fc6119dbc5fcf5c4cc0e1536394b52d73b7904d033836c73db/xxhash-4.0.1-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:f3e1a44af01b6692de0ec6caba5f0bf93ceb36896e02b7fc00952c6ea7ef39e1", size = 330270, upload-time = "2026-08-17T08:20:51.128Z" }, + { url = "https://files.pythonhosted.org/packages/72/9e/e2ab12d40921f3f34c9317637d65e011aeababf8288356ea8d527de2c1d0/xxhash-4.0.1-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:c6fc415b5568bd9accc7187f1729a99707330c0a67a8b9f93c1149ed573ed75d", size = 478555, upload-time = "2026-08-17T08:22:04.183Z" }, + { url = "https://files.pythonhosted.org/packages/6d/32/c6148d39a49efa95f39b4cf0d41ef35a487f3b30f6fb1fc8fe8d8eab577e/xxhash-4.0.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:96d8de55029d42251945531f6aa7590c32b48163c66a43bf29d8657d7446a377", size = 258174, upload-time = "2026-08-17T08:35:21.18Z" }, + { url = "https://files.pythonhosted.org/packages/8f/fb/0b04b68d6c5bc71c7a2c344f1287327b67e607f28fbcfd937697caca64b6/xxhash-4.0.1-cp312-cp312-pyemscripten_2024_0_wasm32.whl", hash = "sha256:0163b5d259de23ae9e07b7eabf435ce4704f6f205589a2b154e6af4be985ce1b", size = 20767, upload-time = "2026-08-17T08:21:00.806Z" }, + { url = "https://files.pythonhosted.org/packages/a6/be/476092aba34d1fcd313e1613a3bb3bc692f253d167b54bc90049043b5034/xxhash-4.0.1-cp312-cp312-win32.whl", hash = "sha256:1216f7ba5683f17a89eb7dcb4bc50a0b743dfe1902278d7b3d0786f538118433", size = 34669, upload-time = "2026-08-17T08:21:49.486Z" }, + { url = "https://files.pythonhosted.org/packages/aa/02/f9413d94fae43cec6d1a74c4f12156c6f4a7f5fd50e1d34defebdee3dec9/xxhash-4.0.1-cp312-cp312-win_amd64.whl", hash = "sha256:5c2d525a3afabcd8e3549d85fc7e111fde6bc302d06a1893fe73adb79823415e", size = 37073, upload-time = "2026-08-17T08:22:04.886Z" }, + { url = "https://files.pythonhosted.org/packages/c1/83/6fe93c1b95acf962bc61a246df09dc2dcce895ccfc1080c9f48d0b652b92/xxhash-4.0.1-cp312-cp312-win_arm64.whl", hash = "sha256:86b2b12bec60c678ed8f5cca0258ad93a8928ebddb6ca7732f0875afe1451d1a", size = 33299, upload-time = "2026-08-17T08:35:12.708Z" }, + { url = "https://files.pythonhosted.org/packages/86/79/9127ff42a887a348dc4ce3211cf1a962836887adee6f57078132bfba78b4/xxhash-4.0.1-graalpy312-graalpy250_312_native-macosx_10_13_x86_64.whl", hash = "sha256:ff48915bf1871a1f19f74c11834c6329443d306cedc0c05fe7fe617810422a80", size = 31836, upload-time = "2026-08-17T08:36:28.261Z" }, + { url = "https://files.pythonhosted.org/packages/0a/e6/f238693bfdd642adb59c99683964d46d9947fe721ff44d3bd850ae675407/xxhash-4.0.1-graalpy312-graalpy250_312_native-macosx_11_0_arm64.whl", hash = "sha256:4a76345f5aceb4ec404918edf9c7f2b5507db864dc0d7455982009ac0890b57b", size = 34453, upload-time = "2026-08-17T08:23:49.795Z" }, + { url = "https://files.pythonhosted.org/packages/40/4b/796ace33cdfb75c91ba6d11615c3bd436355b9f3103e05865bbee9abce57/xxhash-4.0.1-graalpy312-graalpy250_312_native-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:31d86f9e81f3e84e00131ac7c54caf5119ae4ddd82c09c31cff597c813ce1ee2", size = 38488, upload-time = "2026-08-17T08:23:59.901Z" }, + { url = "https://files.pythonhosted.org/packages/ad/23/2d549e5d5d7759eaf9ac2d2d2ab81ff60f1bb2b52cdaae8e5ec5c6524354/xxhash-4.0.1-graalpy312-graalpy250_312_native-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:deca2a30d983d240b8375ec2ee0a4288e72042827fc61df2f7671f8467e4cb2f", size = 38206, upload-time = "2026-08-17T08:36:32.193Z" }, + { url = "https://files.pythonhosted.org/packages/79/98/1ee576b27f78e6107ee4ea8ac03e8a52888dff256e57d560f8282c195563/xxhash-4.0.1-graalpy312-graalpy250_312_native-win_amd64.whl", hash = "sha256:7c343ee174d417a44d0c3355602c0cbbfa52a04d1bbbf1723378c7d2c8f60626", size = 37127, upload-time = "2026-08-17T08:23:42.705Z" }, +] + +[[package]] +name = "yarl" +version = "1.24.5" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "idna" }, + { name = "multidict" }, + { name = "propcache" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/31/33/ebe9e3d1f86c7a0b51094c0a146392045ca1631d2664889539dec8088a33/yarl-1.24.5.tar.gz", hash = "sha256:e81b83143bee16329c23db3c1b2d82b29892fcbcb849186d2f6e98a5abe9a57f", size = 228679, upload-time = "2026-07-20T02:07:45.435Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/1b/84/71d051c850b5af41d168c679d9eb67eb7c55283ac4ee131673edf134bc4e/yarl-1.24.5-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:d693396e5aea78db03decd60aec9ece16c9b40ba00a587f089615ff4e718a81d", size = 136035, upload-time = "2026-07-20T02:05:25.489Z" }, + { url = "https://files.pythonhosted.org/packages/03/4d/8ad27f9a1b7e69313cca5d695b925b48efe51208d3490e0844bae97cabc0/yarl-1.24.5-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:3363fcc96e665878946ad7a106b9a13eac0541766a690ef287c0232ac768b6ec", size = 97642, upload-time = "2026-07-20T02:05:27.429Z" }, + { url = "https://files.pythonhosted.org/packages/ea/b4/05b4131c407006cd1e410e9c6539f16a0945724677e5364447313c15ea3e/yarl-1.24.5-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:9d399bdcfb4a0f659b9b3788bbc89babe63d9a6a65aacdf4d4e7065ff2e6316c", size = 97323, upload-time = "2026-07-20T02:05:29.441Z" }, + { url = "https://files.pythonhosted.org/packages/20/16/e618c875c73e0e39611f20a581b3d5e8d59b8857bf001bee3263044c6deb/yarl-1.24.5-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:90333fd89b43c0d08ac85f3f1447593fc2c66de18c3d6378d7125ea118dc7a54", size = 107741, upload-time = "2026-07-20T02:05:31.367Z" }, + { url = "https://files.pythonhosted.org/packages/d9/9a/c4defeaf3ed33fcb346aacf9c6e971a8d4e2bde04a0310e79abb208e7965/yarl-1.24.5-cp312-cp312-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:665b0a2c463cc9423dd647e0bfd9f4ccc9b50f768c55304d5e9f80b177c1de12", size = 103570, upload-time = "2026-07-20T02:05:33.303Z" }, + { url = "https://files.pythonhosted.org/packages/5f/e7/0e0e0de5865ebd5914537ef486f36c727a59865c3ac0cf5ff1b32aececbf/yarl-1.24.5-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:e006d3a974c4ee19512e5f058abedb6eef36a5e553c14812bdeba1758d812e6d", size = 115815, upload-time = "2026-07-20T02:05:35.292Z" }, + { url = "https://files.pythonhosted.org/packages/2b/27/ca56b700cb170aba25a3893b75355b213935657dc5714d2383354a270e62/yarl-1.24.5-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:e7d42c531243450ef0d4d9c172e7ed6ef052640f195629065041b5add4e058d1", size = 116025, upload-time = "2026-07-20T02:05:37.503Z" }, + { url = "https://files.pythonhosted.org/packages/d6/d0/d56c859b8222116f5d68459199f48359e0bf121b6f65a69bf329b3602ba0/yarl-1.24.5-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f08c7513ecef5aad65687bfdf6bc601ae9fccd04a42904501f8f7141abad9eb9", size = 109835, upload-time = "2026-07-20T02:05:39.506Z" }, + { url = "https://files.pythonhosted.org/packages/70/a2/3a35557e4d1a79425040eba202ccaf08bdc8717680fc77e2498a1ad2e0a5/yarl-1.24.5-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:6c95b17fe34ed802f17e205112e6e10db92275c34fee290aa9bdc55a9c724027", size = 108884, upload-time = "2026-07-20T02:05:41.584Z" }, + { url = "https://files.pythonhosted.org/packages/e4/35/ef4c26356b7913c68983bac2d72a4212b3347af551cb8d250b99b5ed7b7f/yarl-1.24.5-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:56b149b22de33b23b0c6077ab9518c6dcb538ad462e1830e68d06591ccf6e38b", size = 107308, upload-time = "2026-07-20T02:05:43.697Z" }, + { url = "https://files.pythonhosted.org/packages/d5/91/ff0dc66c2ccf3e0153ab97ff61eabab4400e6a5264af427ab30cd69f1857/yarl-1.24.5-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:a8fe66b8f300da93798025a785a5b90b42f3810dc2b72283ff84a41aaaebc293", size = 103646, upload-time = "2026-07-20T02:05:45.895Z" }, + { url = "https://files.pythonhosted.org/packages/74/f0/33b9271c7f881766359d58266fa0811d2e5210ed860e28da7dc6d7786344/yarl-1.24.5-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:377fe3732edbaf78ee74efdf2c9f49f6e99f20e7f9d2649fda3eb4badd77d76e", size = 115305, upload-time = "2026-07-20T02:05:47.832Z" }, + { url = "https://files.pythonhosted.org/packages/ef/65/fd79fb1868c4a80db8661091de525bf430f63c3bea1b20e8b6a84fc7d359/yarl-1.24.5-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:e8ffa78582120024f476a611d7befc123cee59e47e8309d470cf667d806e613b", size = 108404, upload-time = "2026-07-20T02:05:49.604Z" }, + { url = "https://files.pythonhosted.org/packages/ff/ba/dbabe6b262f17a816c70cfc09558dbf03ece3ec76684d02f911a3d3a189c/yarl-1.24.5-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:daba5e594f06114e37db186efd2dd916609071e59daca901a0a2e71f02b142ce", size = 115940, upload-time = "2026-07-20T02:05:51.741Z" }, + { url = "https://files.pythonhosted.org/packages/a5/43/fab2d1dad9d340a268cdde63756a123d069723efff6a372d123fa74a9517/yarl-1.24.5-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:65be18ec59496c13908f02a2472751d9ef840b4f3fb5726f129306bf6a2a7bba", size = 110006, upload-time = "2026-07-20T02:05:53.554Z" }, + { url = "https://files.pythonhosted.org/packages/c4/27/41eb51bbd1b8d89546b83897cfb0164f1e109304fd408dbb151b639eec0f/yarl-1.24.5-cp312-cp312-win_amd64.whl", hash = "sha256:a929d878fec099030c292803b31e5d5540a7b6a31e6a3cc76cb4685fc2a2f51b", size = 97618, upload-time = "2026-07-20T02:05:55.57Z" }, + { url = "https://files.pythonhosted.org/packages/3c/25/b2553764b3d65db711d8f45416351ec4f420847558eb669edcbcaadf5780/yarl-1.24.5-cp312-cp312-win_arm64.whl", hash = "sha256:7ce27823052e2013b597e0c738b13e7e36b8ccb9400df8959417b052ab0fd92c", size = 93018, upload-time = "2026-07-20T02:05:57.554Z" }, + { url = "https://files.pythonhosted.org/packages/61/02/962c1cbfc401a30c1d034dc67ff395f64b52302c6d62de556c1fca99acc0/yarl-1.24.5-py3-none-any.whl", hash = "sha256:a33700d13d9b7d84fd10947b09ff69fb9a792e519c8cb9764a3ca70baa6c23a7", size = 58612, upload-time = "2026-07-20T02:07:43.461Z" }, +] + +[[package]] +name = "zstandard" +version = "0.25.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/fd/aa/3e0508d5a5dd96529cdc5a97011299056e14c6505b678fd58938792794b1/zstandard-0.25.0.tar.gz", hash = "sha256:7713e1179d162cf5c7906da876ec2ccb9c3a9dcbdffef0cc7f70c3667a205f0b", size = 711513, upload-time = "2025-09-14T22:15:54.002Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/82/fc/f26eb6ef91ae723a03e16eddb198abcfce2bc5a42e224d44cc8b6765e57e/zstandard-0.25.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:7b3c3a3ab9daa3eed242d6ecceead93aebbb8f5f84318d82cee643e019c4b73b", size = 795738, upload-time = "2025-09-14T22:16:56.237Z" }, + { url = "https://files.pythonhosted.org/packages/aa/1c/d920d64b22f8dd028a8b90e2d756e431a5d86194caa78e3819c7bf53b4b3/zstandard-0.25.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:913cbd31a400febff93b564a23e17c3ed2d56c064006f54efec210d586171c00", size = 640436, upload-time = "2025-09-14T22:16:57.774Z" }, + { url = "https://files.pythonhosted.org/packages/53/6c/288c3f0bd9fcfe9ca41e2c2fbfd17b2097f6af57b62a81161941f09afa76/zstandard-0.25.0-cp312-cp312-manylinux2010_i686.manylinux2014_i686.manylinux_2_12_i686.manylinux_2_17_i686.whl", hash = "sha256:011d388c76b11a0c165374ce660ce2c8efa8e5d87f34996aa80f9c0816698b64", size = 5343019, upload-time = "2025-09-14T22:16:59.302Z" }, + { url = "https://files.pythonhosted.org/packages/1e/15/efef5a2f204a64bdb5571e6161d49f7ef0fffdbca953a615efbec045f60f/zstandard-0.25.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:6dffecc361d079bb48d7caef5d673c88c8988d3d33fb74ab95b7ee6da42652ea", size = 5063012, upload-time = "2025-09-14T22:17:01.156Z" }, + { url = "https://files.pythonhosted.org/packages/b7/37/a6ce629ffdb43959e92e87ebdaeebb5ac81c944b6a75c9c47e300f85abdf/zstandard-0.25.0-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:7149623bba7fdf7e7f24312953bcf73cae103db8cae49f8154dd1eadc8a29ecb", size = 5394148, upload-time = "2025-09-14T22:17:03.091Z" }, + { url = "https://files.pythonhosted.org/packages/e3/79/2bf870b3abeb5c070fe2d670a5a8d1057a8270f125ef7676d29ea900f496/zstandard-0.25.0-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:6a573a35693e03cf1d67799fd01b50ff578515a8aeadd4595d2a7fa9f3ec002a", size = 5451652, upload-time = "2025-09-14T22:17:04.979Z" }, + { url = "https://files.pythonhosted.org/packages/53/60/7be26e610767316c028a2cbedb9a3beabdbe33e2182c373f71a1c0b88f36/zstandard-0.25.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5a56ba0db2d244117ed744dfa8f6f5b366e14148e00de44723413b2f3938a902", size = 5546993, upload-time = "2025-09-14T22:17:06.781Z" }, + { url = "https://files.pythonhosted.org/packages/85/c7/3483ad9ff0662623f3648479b0380d2de5510abf00990468c286c6b04017/zstandard-0.25.0-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:10ef2a79ab8e2974e2075fb984e5b9806c64134810fac21576f0668e7ea19f8f", size = 5046806, upload-time = "2025-09-14T22:17:08.415Z" }, + { url = "https://files.pythonhosted.org/packages/08/b3/206883dd25b8d1591a1caa44b54c2aad84badccf2f1de9e2d60a446f9a25/zstandard-0.25.0-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:aaf21ba8fb76d102b696781bddaa0954b782536446083ae3fdaa6f16b25a1c4b", size = 5576659, upload-time = "2025-09-14T22:17:10.164Z" }, + { url = "https://files.pythonhosted.org/packages/9d/31/76c0779101453e6c117b0ff22565865c54f48f8bd807df2b00c2c404b8e0/zstandard-0.25.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:1869da9571d5e94a85a5e8d57e4e8807b175c9e4a6294e3b66fa4efb074d90f6", size = 4953933, upload-time = "2025-09-14T22:17:11.857Z" }, + { url = "https://files.pythonhosted.org/packages/18/e1/97680c664a1bf9a247a280a053d98e251424af51f1b196c6d52f117c9720/zstandard-0.25.0-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:809c5bcb2c67cd0ed81e9229d227d4ca28f82d0f778fc5fea624a9def3963f91", size = 5268008, upload-time = "2025-09-14T22:17:13.627Z" }, + { url = "https://files.pythonhosted.org/packages/1e/73/316e4010de585ac798e154e88fd81bb16afc5c5cb1a72eeb16dd37e8024a/zstandard-0.25.0-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:f27662e4f7dbf9f9c12391cb37b4c4c3cb90ffbd3b1fb9284dadbbb8935fa708", size = 5433517, upload-time = "2025-09-14T22:17:16.103Z" }, + { url = "https://files.pythonhosted.org/packages/5b/60/dd0f8cfa8129c5a0ce3ea6b7f70be5b33d2618013a161e1ff26c2b39787c/zstandard-0.25.0-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:99c0c846e6e61718715a3c9437ccc625de26593fea60189567f0118dc9db7512", size = 5814292, upload-time = "2025-09-14T22:17:17.827Z" }, + { url = "https://files.pythonhosted.org/packages/fc/5f/75aafd4b9d11b5407b641b8e41a57864097663699f23e9ad4dbb91dc6bfe/zstandard-0.25.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:474d2596a2dbc241a556e965fb76002c1ce655445e4e3bf38e5477d413165ffa", size = 5360237, upload-time = "2025-09-14T22:17:19.954Z" }, + { url = "https://files.pythonhosted.org/packages/ff/8d/0309daffea4fcac7981021dbf21cdb2e3427a9e76bafbcdbdf5392ff99a4/zstandard-0.25.0-cp312-cp312-win32.whl", hash = "sha256:23ebc8f17a03133b4426bcc04aabd68f8236eb78c3760f12783385171b0fd8bd", size = 436922, upload-time = "2025-09-14T22:17:24.398Z" }, + { url = "https://files.pythonhosted.org/packages/79/3b/fa54d9015f945330510cb5d0b0501e8253c127cca7ebe8ba46a965df18c5/zstandard-0.25.0-cp312-cp312-win_amd64.whl", hash = "sha256:ffef5a74088f1e09947aecf91011136665152e0b4b359c42be3373897fb39b01", size = 506276, upload-time = "2025-09-14T22:17:21.429Z" }, + { url = "https://files.pythonhosted.org/packages/ea/6b/8b51697e5319b1f9ac71087b0af9a40d8a6288ff8025c36486e0c12abcc4/zstandard-0.25.0-cp312-cp312-win_arm64.whl", hash = "sha256:181eb40e0b6a29b3cd2849f825e0fa34397f649170673d385f3598ae17cca2e9", size = 462679, upload-time = "2025-09-14T22:17:23.147Z" }, +] diff --git a/native/rust/dimos-module/Cargo.toml b/native/rust/dimos-module/Cargo.toml index c09298e617..45adc49c1a 100644 --- a/native/rust/dimos-module/Cargo.toml +++ b/native/rust/dimos-module/Cargo.toml @@ -6,9 +6,9 @@ description = "Rust native module SDK for dimos NativeModule framework" license = "Apache-2.0" [dependencies] -dimos-lcm = { git = "https://github.com/dimensionalOS/dimos-lcm.git", branch = "rust-codegen" } +dimos-lcm = { git = "https://github.com/dimensionalOS/dimos-lcm.git", rev = "dd2159513ebfaa7ebc5fc32bf60209cd09aa1ca5" } dimos-module-macros = { version = "=0.1.0", path = "../dimos-module-macros" } -lcm-msgs = { git = "https://github.com/dimensionalOS/dimos-lcm.git", branch = "rust-codegen" } +lcm-msgs = { git = "https://github.com/dimensionalOS/dimos-lcm.git", rev = "dd2159513ebfaa7ebc5fc32bf60209cd09aa1ca5" } nalgebra = "0.35.0" rayon = "1" tokio = { version = "1", features = ["rt-multi-thread", "macros", "sync", "time", "signal", "io-std", "io-util"] } diff --git a/pyproject.toml b/pyproject.toml index 399ded9525..8d5ae39b2a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -100,7 +100,7 @@ classifiers = [ dependencies = [ # Transport Protocols - "dimos-lcm>=0.1.3", + "dimos-lcm @ git+https://github.com/dimensionalOS/dimos-lcm.git@dd2159513ebfaa7ebc5fc32bf60209cd09aa1ca5", "eclipse-zenoh>=1.10.1,<2.0", "PyTurboJPEG==1.8.2", "imagecodecs>=2024.6.1", # JPEG-XL for CompressedImage @@ -206,6 +206,7 @@ visualization = [ learning = [ # dimos.imitation.dataprep dataset writers (lazy-imported per format) + "mcap>=1.2", # read native collection artifacts directly "pyarrow", # LeRobot v3.0 data/episodes parquet "pandas", # LeRobot v3.0 tasks.parquet (task-indexed) "h5py", # HDF5 writer diff --git a/uv.lock b/uv.lock index b28700b2ea..a791da9981 100644 --- a/uv.lock +++ b/uv.lock @@ -1920,6 +1920,7 @@ drone = [ ] learning = [ { name = "h5py" }, + { name = "mcap" }, { name = "pandas", version = "2.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "pandas", version = "3.0.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "pyarrow" }, @@ -2360,7 +2361,7 @@ requires-dist = [ { name = "dimos", extras = ["control"], marker = "extra == 'planning'" }, { name = "dimos", extras = ["planning", "base", "sim", "cpu"], marker = "extra == 'manipulation'" }, { name = "dimos", extras = ["unitree"], marker = "extra == 'unitree-dds'" }, - { name = "dimos-lcm", specifier = ">=0.1.3" }, + { name = "dimos-lcm", git = "https://github.com/dimensionalOS/dimos-lcm.git?rev=dd2159513ebfaa7ebc5fc32bf60209cd09aa1ca5" }, { name = "dimos-viewer", specifier = "==0.32.0a2" }, { name = "dimos-viewer", marker = "extra == 'visualization'", specifier = "==0.32.0a2" }, { name = "drake", marker = "platform_machine != 'aarch64' and sys_platform == 'darwin' and extra == 'planning'", specifier = "==1.45.0" }, @@ -2395,6 +2396,7 @@ requires-dist = [ { name = "lz4", specifier = ">=4.4.5" }, { name = "manifold3d", marker = "extra == 'apriltag'", specifier = ">=2.5.0" }, { name = "matplotlib", marker = "extra == 'planning'", specifier = ">=3.7.1" }, + { name = "mcap", marker = "extra == 'learning'", specifier = ">=1.2" }, { name = "mcap", marker = "extra == 'unitree-dds'", specifier = ">=1.2.0" }, { name = "moondream", marker = "extra == 'perception'" }, { name = "mujoco", marker = "extra == 'sim'", specifier = ">=3.3.4" }, @@ -2671,18 +2673,14 @@ tests-self-hosted = [ [[package]] name = "dimos-lcm" -version = "0.1.3" -source = { registry = "https://pypi.org/simple" } +version = "0.1.4" +source = { git = "https://github.com/dimensionalOS/dimos-lcm.git?rev=dd2159513ebfaa7ebc5fc32bf60209cd09aa1ca5#dd2159513ebfaa7ebc5fc32bf60209cd09aa1ca5" } dependencies = [ { name = "foxglove-websocket" }, { name = "lcm-dimos-fork" }, { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "numpy", version = "2.3.5", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/09/b2/4cdd2bce665ab633313b5d371e0a63fcae2cf77a934a2a9bf820db88a540/dimos_lcm-0.1.3.tar.gz", hash = "sha256:5f7dbd3055f299823bc0e450c59583ad5a2d093c182deec8a86073853881bb09", size = 405688, upload-time = "2026-06-03T07:20:20.552Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/95/19/2d8babf544993508359922ea59db2280d42efa0d5a80d4d0cae22ce40d67/dimos_lcm-0.1.3-py3-none-any.whl", hash = "sha256:63317225a0b4ab0f05e4b656f3f78ac9ba4e914998da2c689ca839cbbd83d57d", size = 1714087, upload-time = "2026-06-03T07:20:22.689Z" }, -] [[package]] name = "dimos-viewer"