Metadata-Version: 2.4
Name: pyroboframes
Version: 1.1.0
Classifier: Development Status :: 3 - Alpha
Classifier: Programming Language :: Rust
Classifier: Programming Language :: Python :: 3
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: MacOS
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Dist: numpy>=1.24
Requires-Dist: pyarrow>=14
Requires-Dist: pytest>=8 ; extra == 'dev'
Requires-Dist: numpy>=1.24 ; extra == 'dev'
Requires-Dist: maturin>=1.7 ; extra == 'dev'
Requires-Dist: mlx>=0.20 ; extra == 'mlx'
Provides-Extra: dev
Provides-Extra: mlx
License-File: LICENSE
Summary: Early/experimental LeRobot dataloader for Apple Silicon & Linux — FFmpeg decode, NumPy/MLX/PyTorch output (hardware decode, zero-copy MLX & parallel prefetch in progress)
Keywords: robotics,robot-learning,lerobot,mlx,apple-silicon,dataloader,videotoolbox
Author-email: Georgi Mammen Mullassery <mullassery@gmail.com>
Requires-Python: >=3.10
Description-Content-Type: text/markdown; charset=UTF-8; variant=GFM
Project-URL: Homepage, https://github.com/Mullassery/PyRoboFrames
Project-URL: Issues, https://github.com/Mullassery/PyRoboFrames/issues
Project-URL: Repository, https://github.com/Mullassery/PyRoboFrames

# PyRoboFrames

[![PyPI](https://img.shields.io/pypi/v/pyroboframes)](https://pypi.org/project/pyroboframes/)
[![Python](https://img.shields.io/pypi/pyversions/pyroboframes)](https://pypi.org/project/pyroboframes/)
[![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)](./LICENSE)
[![Tests](https://img.shields.io/badge/tests-175%20passing-brightgreen)]()

**Fast ML dataloader for robot learning — LeRobot, RLDS, HDF5, NetCDF, hardware video decode, distributed S3/GCS streaming.**

PyRoboFrames is a **foundation library** — load any robot learning dataset, accelerate video decode with hardware (VideoToolbox/NVDEC), validate data quality, and stream to NumPy/MLX/PyTorch/JAX. The heavy lifting runs in a Rust engine; Python is the ergonomic surface.

**For autonomous driving perception and foundation models, see [PyRoboVision](https://github.com/Mullassery/PyRoboVision).**

---

## Installation

```bash
pip install pyroboframes
```

Or with uv:
```bash
uv add pyroboframes
```

**Requires:** Python ≥ 3.10  
**Prebuilt wheels:** macOS (Apple Silicon), Linux (x86_64)  
**From source:** Rust 1.78+ required

Optional extras (install as needed):

```bash
pip install pyroboframes h5py               # HDF5 datasets
pip install pyroboframes xarray netCDF4     # NetCDF datasets
pip install pyroboframes tensorflow-datasets  # RLDS / Open X-Embodiment
pip install pyroboframes fsspec s3fs        # S3 remote streaming
pip install pyroboframes fsspec gcsfs       # GCS remote streaming
pip install pyroboframes ray                # Ray distributed loading
```

---

## Quick Start

### Load LeRobot Datasets

```python
import pyroboframes as prf

ds = prf.RoboFrameDataset.from_path("/path/to/lerobot_dataset")

loader = ds.loader(
    batch_size=64,
    cameras=["observation.images.top"],
    output="torch",       # or "mlx", "numpy", "jax"
    num_workers=4,
    cache_size=4096,      # LRU frame cache (frames)
    episode_prefetch=True,
)

for batch in loader:
    state  = batch["observation.state"]          # [64, state_dim]
    frames = batch["observation.images.top"]     # [64, H, W, 3]
    action = batch["action"]                     # [64, action_dim]
```

### Proprioceptive-Only (No Video) for 10× Speedup

```python
loader = prf.ProprioceptiveLoader(
    dataset_path="/path/to/lerobot_dataset",
    batch_size=256,
    device="mlx",
)

for batch in loader:
    state  = batch["state"]    # [256, state_dim]
    action = batch["action"]   # [256, action_dim]
```

### Temporal Windows for Sequence Models

```python
loader = ds.loader(
    batch_size=32,
    chunk_size=16,
    delta_timestamps={"observation.state": [-0.2, -0.1, 0.0]},
    output="mlx",
)
for batch in loader:
    seq = batch["observation.state"]  # [32, 3, state_dim]
```

---

## What's New in v1.1

### Video Codec Selection — 40–50% Storage Savings

```python
# Write with HEVC (H.265) instead of the H.264 default
prf.write_lerobot_dataset(
    path="/out/dataset",
    features={"observation.state": state_arr, "action": action_arr},
    episode_lengths=[500, 500],
    video_codec="hevc",   # "h264" | "hevc" | "av1"
    video_crf=23,         # lower = better quality, larger file
)

# Standalone video encoding
prf.encode_video_frames(frames, "output.mp4", codec="av1", crf=30)
```

### Data Validation Toolkit

```python
from pyroboframes import DatasetValidator

validator = DatasetValidator(
    ds,
    check_frames=True,    # frame count vs. metadata
    check_temporal=True,  # timestamp gap detection
    check_codec=True,     # sample-decode health check
    sample_rate=0.1,      # probe 10% of episodes
)
report = validator.validate()
print(report.summary())
report.raise_if_errors()
```

### Episode-Level Caching for Repeated Epochs

```python
from pyroboframes import EpisodeCache

cache = EpisodeCache(ds, max_episodes=8)

for epoch in range(10):
    for ep_idx in range(ds.num_episodes()):
        ep = cache.get_episode(ep_idx)   # decoded once, cached after
        states = ep["observation.state"]  # [T, D]

cache.prefetch([0, 1, 2, 3])  # background pre-decode
```

### Cross-Dataset Quality Comparison

```python
from pyroboframes import EpisodeScorer, DatasetQualityProfile, CrossDatasetComparator

scorer = EpisodeScorer()
profile_a = DatasetQualityProfile.from_scores("dataset_a", scorer.score_episodes(df_a))
profile_b = DatasetQualityProfile.from_scores("dataset_b", scorer.score_episodes(df_b))

comparator = CrossDatasetComparator(reference=profile_a)
print(comparator.compare(profile_b))            # Cohen's d, percentile overlap
print(comparator.recommend_mixing_ratio(profile_b))  # curriculum mixing weight
```

### HDF5 / NetCDF / RLDS Format Support

```python
# HDF5 (ROBOMIMIC, ACT, custom) — pip install h5py
from pyroboframes import HDF5Dataset, convert_hdf5
convert_hdf5("robomimic.hdf5", "/out/lerobot")

# NetCDF (scientific/simulation datasets) — pip install xarray netCDF4
from pyroboframes import NetCDFDataset, convert_netcdf
convert_netcdf("sim_data.nc", "/out/lerobot", episode_breaks=[0, 500, 1200])

# RLDS / Open X-Embodiment — pip install tensorflow-datasets
from pyroboframes import RLDSDataset, convert_rlds
convert_rlds("fractal20220817_data", "/out/lerobot", split="train")
```

### Remote S3/GCS Streaming + Ray Distributed Loading

```python
# Stream from S3
from pyroboframes import RemoteDataset
ds = RemoteDataset.from_s3("s3://my-bucket/lerobot_dataset").open()
ds.prefetch_episodes([0, 1, 2, 3])   # background download
loader = ds.loader(batch_size=32)

# Ray distributed — pip install ray
from pyroboframes import RayDistributedLoader, shard_episodes
loader = RayDistributedLoader(
    "/path/to/dataset", num_workers=4, rank=0, world_size=4, batch_size=32
)

# Or just shard episodes yourself
my_episodes = shard_episodes(total_episodes=200, world_size=4, rank=0)
# → [0, 4, 8, …, 196]
```

---

## Full Feature Table

| Feature | Status | Notes |
|---|---|---|
| **LeRobot v3.0 loading** | ✅ | Full schema support |
| **Video decode** | ✅ | FFmpeg + VideoToolbox + NVDEC |
| **Proprioceptive loader** | ✅ | 10× speedup (no video) |
| **Temporal windows** | ✅ | Multi-timestep sequences |
| **Multi-camera batching** | ✅ | Arbitrary camera combinations |
| **Output formats** | ✅ | NumPy, MLX, PyTorch, JAX |
| **Parallel prefetch** | ✅ | num_workers for async loading |
| **Data augmentation** | ✅ | Rotate, flip, crop, color jitter |
| **Video codec selection** | ✅ | H.264 / HEVC / AV1 + CRF control |
| **Dataset validation** | ✅ | Temporal gaps, missing frames, codec health |
| **Episode caching** | ✅ | RAM-based LRU cache, background prefetch |
| **MCAP ingestion** | ✅ | JSON, protobuf, CDR support |
| **ROS 2 bag ingestion** | ✅ | .db3 native format |
| **HDF5 ingestion** | ✅ | ROBOMIMIC, ACT, custom layouts |
| **NetCDF ingestion** | ✅ | Scientific/simulation datasets |
| **RLDS / Open X-Embodiment** | ✅ | tensorflow-datasets integration |
| **Episode quality scoring** | ✅ | Diversity, sharpness, state variance |
| **Cross-dataset comparison** | ✅ | Cohen's d, percentile ranking, mixing ratio |
| **S3/GCS streaming** | ✅ | fsspec-backed remote datasets |
| **Ray distributed loading** | ✅ | Episode sharding across Ray workers |
| **Streaming ingestion** | ✅ | Kafka, MQTT real-time data |
| **Distributed loading** | ✅ | Multi-GPU synchronized sampling |

---

## Test Coverage: 175 Tests Passing ✅

```
Dataloader:       30 tests
Video decode:     25 tests
Proprioceptive:   16 tests
Augmentation:     15 tests
Temporal ops:     12 tests
Quality/scoring:  17 tests    (+7 cross-dataset)
Validation:       13 tests
Caching:           5 tests
HDF5:              7 tests
NetCDF:            7 tests
Distributed:       8 tests
Streaming:         7 tests
Codecs:            7 tests    (+3 round-trip)
Other:             6 tests
```

```bash
pytest tests/ -v
```

---

## GPU Support

- **Apple Silicon**: VideoToolbox hardware decode, MLX zero-copy arrays
- **NVIDIA**: NVDEC hardware decode, PyTorch CUDA acceleration
- **CPU**: NumPy fallback (~10× slower than hardware)

```python
loader = ds.loader(device="auto", ...)   # auto-detect
loader = ds.loader(device="mlx", ...)   # Apple Silicon
loader = ds.loader(device="cuda", ...)  # NVIDIA
loader = ds.loader(device="cpu", ...)   # CPU
```

---

## Use Cases

- **LeRobot policy training** — Fast loading for imitation learning
- **Open X-Embodiment fine-tuning** — RLDS ingestion + LeRobot conversion
- **Large-scale cloud training** — S3/GCS streaming + Ray distribution
- **Multi-dataset curriculum** — Cross-dataset quality comparison + mixing ratios
- **Data quality auditing** — Validate integrity before long training runs
- **Legacy dataset migration** — HDF5/NetCDF → LeRobot conversion

---

## Performance

- **Video decode:** 100+ FPS (hardware-accelerated on macOS/CUDA)
- **Dataloader throughput:** 50–100 images/sec (PyTorch, Mac M3)
- **Proprioceptive loader:** 1,000+ batch/sec (no video decode)
- **Storage savings:** 40–50% with HEVC vs H.264 at equivalent quality

---

## Architecture

```
PyRoboFrames (Rust core + Python surface)

Input: LeRobot / HDF5 / NetCDF / RLDS / MCAP / ROS2 / S3 / GCS
   ↓
Format Converters  →  LeRobot v3.0 (Parquet + MP4)
   ↓
Rust Decoder (VideoToolbox / NVDEC / FFmpeg)
   ↓
RoboFrameDataset (episode index, frame manifest)
   ↓
Loader (temporal windows, augmentation, caching, batching)
   ↓
Output: NumPy / MLX / PyTorch / JAX
   ↓
Your training loop
```

### Module Organization

```
pyroboframes/
├── RoboFrameDataset      # Load LeRobot datasets
├── ProprioceptiveLoader  # State/action only (no video)
├── DataLoader            # Flexible batching + augmentation
├── EpisodeCache          # RAM-based episode LRU cache
├── DatasetValidator      # Deep data quality checks
├── hdf5                  # HDF5 reader + converter
├── netcdf                # NetCDF reader + converter
├── rlds                  # RLDS / Open X-Embodiment reader
├── distributed           # RemoteDataset, RayDistributedLoader, shard_episodes
├── quality               # EpisodeScorer, CrossDatasetComparator
├── backend/              # Device abstractions (MLX, PyTorch, JAX)
├── transforms/           # Augmentation pipelines
└── [streaming, sensor_fusion, depth_io, ...]
```

---

## Related Projects

- **[LeRobot](https://github.com/huggingface/lerobot)** — Robot learning datasets
- **[PyRoboVision](https://github.com/Mullassery/PyRoboVision)** — Autonomous driving perception + foundation models
- **[MLX](https://github.com/ml-explore/mlx)** — Apple Silicon ML framework
- **[Open X-Embodiment](https://robotics-transformer-x.github.io/)** — Cross-embodiment robotics datasets

---

## Documentation

- [ARCHITECTURE.md](./ARCHITECTURE.md) — Design and implementation
- [CONTRIBUTING.md](./CONTRIBUTING.md) — How to contribute
- [CHANGELOG.md](./CHANGELOG.md) — Version history
- [IMPLEMENTATION_ROADMAP.md](./IMPLEMENTATION_ROADMAP.md) — Planned features

---

## Community

- **GitHub Issues** — [Ask questions, report bugs](https://github.com/Mullassery/PyRoboFrames/issues)
- **GitHub Discussions** — [Share ideas and best practices](https://github.com/Mullassery/PyRoboFrames/discussions)
- **Code of Conduct** — [Be respectful and constructive](./CODE_OF_CONDUCT.md)

---

## License

[MIT](./LICENSE) © Georgi Mammen Mullassery

---

## Citation

```bibtex
@software{mullassery2025pyroboframes,
  title={PyRoboFrames: Fast ML dataloader for robot learning},
  author={Mullassery, Georgi},
  url={https://github.com/Mullassery/PyRoboFrames},
  year={2025}
}
```

