Metadata-Version: 2.4
Name: traceplane
Version: 0.5.3
Summary: Streaming dataloader for robotics trajectory datasets
Project-URL: Homepage, https://traceplane.ai
Project-URL: Documentation, https://docs.traceplane.ai
Project-URL: Repository, https://github.com/traceplane/traceplane
Project-URL: Issues, https://github.com/traceplane/traceplane/issues
Author-email: Traceplane <hello@traceplane.ai>
License-Expression: Apache-2.0
Keywords: datasets,imitation-learning,lerobot,robotics,trajectories
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: Apache Software License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Requires-Dist: fsspec>=2023.1
Requires-Dist: numpy>=1.24
Requires-Dist: pandas>=2.0
Requires-Dist: pyarrow>=14.0
Requires-Dist: pyyaml>=6.0
Requires-Dist: requests>=2.28
Provides-Extra: all
Requires-Dist: av>=11.0; extra == 'all'
Requires-Dist: fastapi>=0.110; extra == 'all'
Requires-Dist: gcsfs>=2023.1; extra == 'all'
Requires-Dist: huggingface-hub>=0.20; extra == 'all'
Requires-Dist: huggingface-hub>=0.34; extra == 'all'
Requires-Dist: jax>=0.4; extra == 'all'
Requires-Dist: jaxlib>=0.4; extra == 'all'
Requires-Dist: mujoco>=3.0; extra == 'all'
Requires-Dist: numpy>=1.24; extra == 'all'
Requires-Dist: pillow>=10.0; extra == 'all'
Requires-Dist: pydantic>=2.0; extra == 'all'
Requires-Dist: s3fs>=2023.1; extra == 'all'
Requires-Dist: sentence-transformers>=2.0; extra == 'all'
Requires-Dist: tensorflow>=2.14; extra == 'all'
Requires-Dist: torch>=2.0; extra == 'all'
Requires-Dist: transformers>=4.40; extra == 'all'
Requires-Dist: uvicorn>=0.27; extra == 'all'
Provides-Extra: dev
Requires-Dist: av>=11.0; extra == 'dev'
Requires-Dist: fastapi>=0.110; extra == 'dev'
Requires-Dist: httpx>=0.27; extra == 'dev'
Requires-Dist: mujoco>=3.0; extra == 'dev'
Requires-Dist: numpy>=1.24; extra == 'dev'
Requires-Dist: pydantic>=2.0; extra == 'dev'
Requires-Dist: pytest>=7; extra == 'dev'
Requires-Dist: torch>=2.0; extra == 'dev'
Requires-Dist: uvicorn>=0.27; extra == 'dev'
Provides-Extra: embeddings
Requires-Dist: sentence-transformers>=2.0; extra == 'embeddings'
Provides-Extra: gcs
Requires-Dist: gcsfs>=2023.1; extra == 'gcs'
Provides-Extra: hub
Requires-Dist: huggingface-hub>=0.20; extra == 'hub'
Provides-Extra: jax
Requires-Dist: jax>=0.4; extra == 'jax'
Requires-Dist: jaxlib>=0.4; extra == 'jax'
Provides-Extra: mujoco
Requires-Dist: mujoco>=3.0; extra == 'mujoco'
Requires-Dist: numpy>=1.24; extra == 'mujoco'
Provides-Extra: s3
Requires-Dist: s3fs>=2023.1; extra == 's3'
Provides-Extra: server
Requires-Dist: fastapi>=0.110; extra == 'server'
Requires-Dist: pydantic>=2.0; extra == 'server'
Requires-Dist: uvicorn>=0.27; extra == 'server'
Provides-Extra: sim
Requires-Dist: numpy>=1.24; extra == 'sim'
Requires-Dist: torch>=2.0; extra == 'sim'
Provides-Extra: tf
Requires-Dist: tensorflow>=2.14; extra == 'tf'
Provides-Extra: torch
Requires-Dist: torch>=2.0; extra == 'torch'
Provides-Extra: training
Requires-Dist: torch>=2.0; extra == 'training'
Provides-Extra: vision
Requires-Dist: av>=11.0; extra == 'vision'
Requires-Dist: huggingface-hub>=0.34; extra == 'vision'
Requires-Dist: pillow>=10.0; extra == 'vision'
Requires-Dist: torch>=2.0; extra == 'vision'
Requires-Dist: transformers>=4.40; extra == 'vision'
Description-Content-Type: text/markdown

# Traceplane

Python SDK for the Traceplane trajectory data platform.

## Installation

```bash
pip install traceplane
```

With framework extras:

```bash
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[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:

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

## Quick Start

```python
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

```python
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.

```python
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.

```python
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:

```bash
pip install traceplane[vision]
```

```bash
# 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
```

```python
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.

```bash
# 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
```

## API Reference

Full documentation: [docs.traceplane.ai](https://docs.traceplane.ai)

## License

Apache-2.0
