traceplane 0.3.0

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 rollout writer
pip install traceplane[all]       # Everything

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
• 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},
        )

API Reference

Full documentation: docs.traceplane.ai

License

Apache-2.0