traceplane 0.5.5

Streaming dataloader for robotics trajectory datasets

Traceplane

Python SDK for the Traceplane trajectory data platform.

Installation

pip install traceplane

With framework extras:

pip install traceplane[torch]     # PyTorch DataLoader
pip install traceplane[jax]       # JAX support
pip install traceplane[training]  # Diffusion policy training
pip install traceplane[sim]       # Isaac Sim backend (NVIDIA GPU required)
pip install traceplane[mujoco]    # MuJoCo backend (CPU-friendly, works on Blackwell)
pip install traceplane[vlm]       # VLM auto-annotation (Claude)
pip install traceplane[all]       # Everything

Picking a sim backend

Backend	Install	Use when
traceplane[sim] (Isaac Sim)	pip install isaacsim --extra-index-url https://pypi.nvidia.com then pip install traceplane[sim]	Photorealistic rendering, USD assets, NVIDIA GPU available, host is pre-Blackwell
traceplane[mujoco] (MuJoCo)	pip install traceplane[mujoco]	CPU eval, CI runs, Blackwell hosts (RTX 5080+) where Isaac Sim 5.x currently segfaults, no NVIDIA dependency

Both backends share the same config + metrics layer; pick at runtime:

traceplane-sim-eval --sim-backend=mujoco --task=pick_place_cube --num-episodes=10

Quick Start

from traceplane import TraceplaneClient

client = TraceplaneClient("https://api.traceplane.ai", api_key="tp_live_...")

# Register a dataset
client.register("my_data", "/path/to/dataset", include_data=True)

# Query with SQL
rows = client.sql_rows("SELECT * FROM my_data WHERE frame_count > 100")

# Upload data
client.upload_dataset("my_data", "/path/to/parquet/files/")

# Vector search
results = client.search_similar("my_data", episode_index=0, k=5)

Features

• SQL query engine -- register datasets and query with full SQL, including vector UDFs (vec_mean, vec_norm, vec_cosine_sim, etc.)
• Streaming dataloaders -- PyTorch, JAX, and TensorFlow adapters with windowed sampling
• LeRobot format -- native reader and writer for LeRobot v2/v3 datasets (Parquet + MP4)
• Embodiment registry -- 8 built-in robot profiles (Franka, ALOHA, Unitree G1/H1, UR5e, KUKA iiwa, xArm 7, Stretch 3) with joint-limit + action-space validation, shared verbatim with the Rust backend
• Similarity search -- find related episodes via embedding-based vector search; DINOv2 keyframe embeddings for visual similarity (traceplane[vision])
• Dataset upload -- push local Parquet files to the platform
• Retargeting -- XR hand poses to robot action space via calibration bridge
• Training -- built-in diffusion policy training with traceplane-train CLI
• Isaac Sim round-trip -- capture headless rollouts as canonical LeRobot v2 episodes via SimEpisodeWriter

Training Integration

from traceplane import LeRobotReader
from traceplane.torch import TorchEpisodeLoader

reader = LeRobotReader("/path/to/lerobot/dataset")
loader = TorchEpisodeLoader(reader, batch_size=32, window_size=16)

for batch in loader:
    observations = batch["observation"]
    actions = batch["action"]
    # ... your training loop

Embodiment Profiles

Built-in kinematic profiles for 8 robots, each with joint-limit + action-space validation. The same YAML files drive both this SDK and the Rust backend, so cross-embodiment queries and retargeting validation stay consistent end-to-end.

from traceplane.embodiments import get_profile, list_profiles

print(list_profiles())
# ['franka_panda', 'aloha_v2', 'unitree_g1', 'unitree_h1',
#  'ur5e', 'kuka_iiwa', 'xarm7', 'stretch3']

g1 = get_profile("unitree_g1")
print(g1.total_dof())           # 23
print(g1.group_names())         # ['left_arm', 'right_arm', 'waist', ...]
arm_action = g1.project_action(full_action_23d, "left_arm")  # 7-dim subvector

Sim Rollout Round-trip

SimEpisodeWriter turns Isaac Sim rollouts (or any simulator's per-step actions/states) into canonical LeRobot v2 datasets, ready to be re-ingested by Traceplane or consumed by any LeRobot-compatible tool.

from traceplane.embodiments import get_profile
from traceplane.sim.writer import SimEpisodeWriter

profile = get_profile("franka_panda")
with SimEpisodeWriter("./rollouts", profile, fps=30.0, task="pick-cube") as w:
    for ep in rollouts:
        w.add_episode(
            actions=ep.actions,        # (T, profile.action_dim)
            states=ep.joint_positions, # (T, profile.state_dim)
            timestamps=ep.timestamps,  # optional; defaults to 1/fps spacing
            metadata={"success": ep.success, "seed": ep.seed},
        )

Vision Frame Embeddings

Encode keyframes from a dataset's videos with a pretrained DINOv2 model for visual similarity search and near-duplicate detection. This is the vision counterpart to traceplane.embeddings (trajectory statistics + text labels). Requires the optional vision extra:

pip install traceplane[vision]

# Embed every 5th keyframe of each episode's video to a Parquet shard
python -m traceplane.frame_embeddings embed ./my-dataset --stride 5

# Find the 10 nearest frames to a query frame
python -m traceplane.frame_embeddings search \
  ./my-dataset/embeddings/frame_embeddings.parquet \
  --query-episode 0 --query-frame 50 --k 10

from traceplane.frame_embeddings import compute_frame_embeddings, search_frames

path = compute_frame_embeddings("./my-dataset", stride=5)   # -> Parquet
hits = search_frames(path, query_episode=0, query_frame=50, k=10)

Embeddings are the L2-normalised DINOv2 CLS token (768-dim), stored as Parquet and searched with brute-force cosine — consistent with the existing vec_cosine_sim path. (ANN/Lance is deferred until frame counts require it.)

Once embeddings exist, traceplane check --ml adds ML-powered QA to the standard dataset CI report: outlier flagging (k-NN cosine distance, z-scored) and near-duplicate clustering, layered on top of the existing 50+ rule-based checks.

# After running `python -m traceplane.frame_embeddings embed ./my-dataset`:
traceplane check ./my-dataset --ml
# Optional knobs:
traceplane check ./my-dataset --ml --outlier-z 3.0 --dedupe-threshold 0.97
traceplane check ./my-dataset --ml --embeddings /elsewhere/fe.parquet

VLM Auto-Annotation

Most public LeRobot datasets ship with thin language_instruction and no per-object semantic labels. The VLM layer (traceplane[vlm]) samples keyframes from each episode and asks a vision-language model for a structured annotation (task sentence, objects, action verbs, scene). Successful annotations are merge-upserted into meta/annotations.jsonl inside the dataset, keyed by episode_index. Downstream read-back into LeRobotReader.metadata and qa --semantic lands in a follow-up — until then the JSONL is an on-disk artifact for other tools to consume.

pip install traceplane[vlm]
export ANTHROPIC_API_KEY=sk-...

Programmatic use:

from traceplane.vlm import get_vlm
from traceplane.health_report import sample_episode_keyframes

vlm = get_vlm()  # Claude by default; VLM_PROVIDER=gemini reserved for bulk
frames_by_ep = sample_episode_keyframes("./my-dataset", episode_indices=[0, 1, 2])
for ep, frames in frames_by_ep.items():
    ann = vlm.annotate(frames, task_hint="pick_place")
    print(ep, ann["language_instruction"], ann["objects"])

Or via the QA sidecar (same code paths the app uses), with the dataset path served behind a shared token:

traceplane-qa &  # starts the FastAPI sidecar on :8080
curl -X POST http://localhost:8080/annotate \
  -H "X-QA-Token: $TRACEPLANE_QA_TOKEN" \
  -H "content-type: application/json" \
  -d '{"path": "./my-dataset", "keyframes": 6}'

Provider selection is via VLM_PROVIDER (default claude) and VLM_MODEL (default claude-sonnet-4-6; claude-haiku-4-5-20251001 for cheap bulk).

API Reference

Full documentation: docs.traceplane.ai

License

Apache-2.0