traceplane 0.1.0__py3-none-any.whl

traceplane/__init__.py

@@ -0,0 +1,16 @@
+"""Traceplane — streaming dataloader for robotics trajectory datasets."""
+
+from traceplane.dataset import Episode, TrajectoryDataset
+from traceplane.lerobot_reader import LeRobotReader
+from traceplane.windowing import WindowedDataset
+from traceplane.query import TraceplaneClient
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "Episode",
+    "TrajectoryDataset",
+    "LeRobotReader",
+    "WindowedDataset",
+    "TraceplaneClient",
+]

traceplane/dataset.py

@@ -0,0 +1,93 @@
+"""Core dataset abstractions."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Iterator, Sequence
+
+import numpy as np
+
+
+@dataclass
+class Episode:
+    """A single trajectory episode.
+
+    Attributes:
+        episode_id: Unique identifier (e.g. "episode_000042").
+        observations: Dict of observation arrays keyed by modality.
+            Common keys: "state" (proprioception), "action", image camera names.
+            Each value is shape (T, D) for vectors or (T, H, W, C) for images.
+        actions: Action array, shape (T, action_dim).
+        timestamps: Monotonic timestamps in seconds, shape (T,).
+        metadata: Arbitrary episode metadata (task label, fps, success, etc.).
+    """
+
+    episode_id: str
+    observations: dict[str, np.ndarray] = field(default_factory=dict)
+    actions: np.ndarray = field(default_factory=lambda: np.empty((0,)))
+    timestamps: np.ndarray = field(default_factory=lambda: np.empty((0,)))
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+    @property
+    def length(self) -> int:
+        """Number of timesteps."""
+        if self.actions.ndim >= 1 and self.actions.shape[0] > 0:
+            return self.actions.shape[0]
+        if self.timestamps.ndim >= 1 and self.timestamps.shape[0] > 0:
+            return self.timestamps.shape[0]
+        for v in self.observations.values():
+            if hasattr(v, "shape") and v.shape[0] > 0:
+                return v.shape[0]
+        return 0
+
+    @property
+    def action_dim(self) -> int:
+        if self.actions.ndim == 2:
+            return self.actions.shape[1]
+        return 0
+
+
+class TrajectoryDataset:
+    """Abstract base for trajectory datasets.
+
+    Subclasses must implement ``__len__`` and ``__getitem__``.
+    """
+
+    def __len__(self) -> int:
+        raise NotImplementedError
+
+    def __getitem__(self, idx: int) -> Episode:
+        raise NotImplementedError
+
+    def __iter__(self) -> Iterator[Episode]:
+        for i in range(len(self)):
+            yield self[i]
+
+    def episode_ids(self) -> list[str]:
+        """Return all episode IDs in order."""
+        raise NotImplementedError
+
+    def filter(self, episode_ids: Sequence[str]) -> "FilteredDataset":
+        """Return a view containing only the specified episodes."""
+        return FilteredDataset(self, list(episode_ids))
+
+
+class FilteredDataset(TrajectoryDataset):
+    """A filtered view over another dataset."""
+
+    def __init__(self, parent: TrajectoryDataset, episode_ids: list[str]):
+        self._parent = parent
+        self._ids = episode_ids
+        # Build index map: episode_id -> parent index
+        parent_ids = parent.episode_ids()
+        self._id_to_idx = {eid: i for i, eid in enumerate(parent_ids)}
+        self._indices = [self._id_to_idx[eid] for eid in episode_ids if eid in self._id_to_idx]
+
+    def __len__(self) -> int:
+        return len(self._indices)
+
+    def __getitem__(self, idx: int) -> Episode:
+        return self._parent[self._indices[idx]]
+
+    def episode_ids(self) -> list[str]:
+        return [self._ids[i] for i in range(len(self._indices))]

traceplane/embeddings.py

@@ -0,0 +1,207 @@
+"""Episode embedding computation for similarity search.
+
+Computes per-episode feature vectors and writes them to Parquet for
+use with the DataFusion ``vec_cosine_sim`` UDF.
+
+Two strategies:
+    - **Trajectory features** (always available): statistical aggregates
+      of observation.state and action vectors per episode.
+    - **Text embeddings** (optional, requires ``sentence-transformers``):
+      encodes task_label strings with a pretrained language model.
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import pyarrow as pa
+import pyarrow.parquet as pq
+
+from traceplane.dataset import Episode
+from traceplane.lerobot_reader import LeRobotReader
+
+
+def compute_trajectory_embedding(episode: Episode) -> np.ndarray:
+    """Compute a feature vector from an episode's state and action statistics.
+
+    Concatenates [mean, std, min, max] per dimension for both
+    observation.state and action, then L2-normalises.
+
+    Returns:
+        1-D float32 array.
+    """
+    parts: list[np.ndarray] = []
+
+    # State features
+    state = episode.observations.get("state")
+    if state is not None and state.ndim == 2 and state.shape[0] > 0:
+        parts.extend(_stat_features(state))
+
+    # Action features
+    if episode.actions.ndim == 2 and episode.actions.shape[0] > 0:
+        parts.extend(_stat_features(episode.actions))
+
+    if not parts:
+        # Fallback: zero vector
+        return np.zeros(1, dtype=np.float32)
+
+    vec = np.concatenate(parts).astype(np.float32)
+
+    # L2 normalise
+    norm = np.linalg.norm(vec)
+    if norm > 1e-8:
+        vec /= norm
+
+    return vec
+
+
+def _stat_features(arr: np.ndarray) -> list[np.ndarray]:
+    """Compute [mean, std, min, max] per dimension."""
+    return [
+        arr.mean(axis=0),
+        arr.std(axis=0),
+        arr.min(axis=0),
+        arr.max(axis=0),
+    ]
+
+
+def compute_text_embeddings(
+    labels: list[str],
+    model_name: str = "sentence-transformers/all-MiniLM-L6-v2",
+) -> np.ndarray:
+    """Encode text labels with a sentence-transformer model.
+
+    Returns:
+        (N, D) float32 array of embeddings.
+    """
+    try:
+        from sentence_transformers import SentenceTransformer
+    except ImportError:
+        raise ImportError(
+            "sentence-transformers is required for text embeddings. "
+            "Install with: pip install traceplane[embeddings]"
+        )
+    model = SentenceTransformer(model_name)
+    embeddings = model.encode(labels, show_progress_bar=False, normalize_embeddings=True)
+    return np.asarray(embeddings, dtype=np.float32)
+
+
+def compute_embeddings(
+    dataset_path: str,
+    output_dir: str | None = None,
+    include_text: bool = False,
+    episode_indices: list[int] | None = None,
+    storage_options: dict[str, Any] | None = None,
+) -> str:
+    """Compute episode embeddings and write to Parquet.
+
+    Args:
+        dataset_path: Path to a LeRobot dataset.
+        output_dir: Output directory. Defaults to ``{dataset_path}/embeddings``.
+        include_text: Also compute text embeddings from task labels.
+        episode_indices: Subset of episodes. None = all.
+        storage_options: fsspec options for remote datasets.
+
+    Returns:
+        Path to the written Parquet file.
+    """
+    reader = LeRobotReader(
+        dataset_path,
+        storage_options=storage_options,
+        episode_indices=episode_indices,
+    )
+
+    n = len(reader)
+    if n == 0:
+        raise ValueError("Dataset has no episodes")
+
+    print(f"Computing embeddings for {n} episodes...", file=sys.stderr)
+
+    indices: list[int] = []
+    labels: list[str] = []
+    traj_embeddings: list[np.ndarray] = []
+
+    for i in range(n):
+        ep = reader[i]
+        traj_emb = compute_trajectory_embedding(ep)
+        traj_embeddings.append(traj_emb)
+
+        # Extract episode index from metadata or ID
+        ep_idx = ep.metadata.get("episode_index", i)
+        if isinstance(ep_idx, str):
+            try:
+                ep_idx = int(ep_idx.split("_")[-1])
+            except (ValueError, IndexError):
+                ep_idx = i
+        indices.append(int(ep_idx))
+        labels.append(ep.metadata.get("task_label", ""))
+
+        if (i + 1) % 50 == 0 or i == n - 1:
+            print(f"  {i + 1}/{n}", file=sys.stderr)
+
+    # Build arrow arrays
+    traj_dim = traj_embeddings[0].shape[0]
+    arrow_traj = pa.list_(pa.float32())
+
+    columns: dict[str, Any] = {
+        "episode_index": pa.array(indices, type=pa.int64()),
+        "task_label": pa.array(labels, type=pa.utf8()),
+        "trajectory_embedding": pa.array(
+            [emb.tolist() for emb in traj_embeddings],
+            type=arrow_traj,
+        ),
+    }
+
+    # Optional text embeddings
+    if include_text:
+        unique_labels = list(set(labels))
+        print(f"Computing text embeddings for {len(unique_labels)} unique labels...", file=sys.stderr)
+        text_embs = compute_text_embeddings(unique_labels)
+        label_to_emb = {lbl: text_embs[i] for i, lbl in enumerate(unique_labels)}
+        text_vecs = [label_to_emb[lbl].tolist() for lbl in labels]
+        columns["text_embedding"] = pa.array(
+            text_vecs,
+            type=pa.list_(pa.float32()),
+        )
+
+    table = pa.table(columns)
+
+    # Write
+    if output_dir is None:
+        output_dir = os.path.join(dataset_path, "embeddings")
+    os.makedirs(output_dir, exist_ok=True)
+    output_path = os.path.join(output_dir, "episode_embeddings.parquet")
+    pq.write_table(table, output_path)
+
+    print(f"Written {n} embeddings ({traj_dim}-dim) to {output_path}", file=sys.stderr)
+    return output_path
+
+
+def main(argv: list[str] | None = None) -> None:
+    """CLI entry point for embedding computation."""
+    parser = argparse.ArgumentParser(
+        description="Compute episode embeddings for similarity search",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument("dataset_path", help="Path to LeRobot dataset")
+    parser.add_argument("--output-dir", help="Output directory (default: {dataset_path}/embeddings)")
+    parser.add_argument("--include-text", action="store_true", help="Compute text embeddings (requires sentence-transformers)")
+    parser.add_argument("--episodes", type=int, nargs="+", help="Specific episode indices")
+
+    args = parser.parse_args(argv)
+    path = compute_embeddings(
+        args.dataset_path,
+        output_dir=args.output_dir,
+        include_text=args.include_text,
+        episode_indices=args.episodes,
+    )
+    print(path)
+
+
+if __name__ == "__main__":
+    main()

traceplane/jax.py

@@ -0,0 +1,108 @@
+"""JAX integration — convert episodes and windows to JAX arrays."""
+
+from __future__ import annotations
+
+from typing import Any, Iterator
+
+import numpy as np
+
+try:
+    import jax
+    import jax.numpy as jnp
+except ImportError:
+    raise ImportError(
+        "JAX is required for traceplane.jax. "
+        "Install it with: pip install traceplane[jax]"
+    )
+
+from traceplane.dataset import TrajectoryDataset
+from traceplane.windowing import WindowedDataset
+
+
+def episodes_to_jax(
+    dataset: TrajectoryDataset,
+) -> list[dict[str, Any]]:
+    """Convert all episodes to dicts of JAX arrays.
+
+    Returns a list of dicts, each containing:
+        - ``actions``: jax array (T, action_dim)
+        - ``timestamps``: jax array (T,)
+        - observation keys: jax arrays
+        - ``episode_id``: str
+        - ``metadata``: dict
+    """
+    results = []
+    for ep in dataset:
+        item: dict[str, Any] = {
+            "episode_id": ep.episode_id,
+            "metadata": ep.metadata,
+            "actions": jnp.array(ep.actions, dtype=jnp.float32),
+            "timestamps": jnp.array(ep.timestamps, dtype=jnp.float64),
+        }
+        for key, arr in ep.observations.items():
+            item[key] = jnp.array(arr, dtype=jnp.float32)
+        results.append(item)
+    return results
+
+
+def windowed_iterator(
+    dataset: TrajectoryDataset,
+    obs_horizon: int = 2,
+    action_horizon: int = 16,
+    stride: int = 1,
+    batch_size: int = 64,
+    shuffle: bool = True,
+    seed: int = 0,
+) -> Iterator[dict[str, Any]]:
+    """Yield batched training windows as dicts of JAX arrays.
+
+    Args:
+        dataset: Source trajectory dataset.
+        obs_horizon: Observation history length.
+        action_horizon: Action prediction horizon.
+        stride: Window stride.
+        batch_size: Batch size.
+        shuffle: Shuffle sample order.
+        seed: Random seed for shuffling.
+
+    Yields:
+        Dicts with ``obs_history`` (dict of jax arrays), ``action_chunk``
+        (jax array), ``episode_ids`` (list[str]), ``timesteps`` (list[int]).
+    """
+    windowed = WindowedDataset(
+        dataset,
+        obs_horizon=obs_horizon,
+        action_horizon=action_horizon,
+        stride=stride,
+    )
+
+    indices = np.arange(len(windowed))
+    if shuffle:
+        rng = np.random.default_rng(seed)
+        rng.shuffle(indices)
+
+    for start in range(0, len(indices), batch_size):
+        batch_indices = indices[start : start + batch_size]
+        windows = [windowed[int(i)] for i in batch_indices]
+
+        # Collate
+        obs_keys = set()
+        for w in windows:
+            obs_keys.update(w.obs_history.keys())
+
+        batch: dict[str, Any] = {
+            "episode_ids": [w.episode_id for w in windows],
+            "timesteps": [w.timestep for w in windows],
+            "action_chunk": jnp.stack(
+                [jnp.array(w.action_chunk, dtype=jnp.float32) for w in windows]
+            ),
+            "obs_history": {},
+        }
+        for key in obs_keys:
+            arrays = [
+                jnp.array(w.obs_history.get(key, np.zeros((obs_horizon, 1))), dtype=jnp.float32)
+                for w in windows
+            ]
+            batch["obs_history"][key] = jnp.stack(arrays)
+
+        yield batch