File size: 4,276 Bytes
319eb16 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 | # MIT-Franka-P-Rank Dataset Guide
This guide explains how to integrate and use the MIT-Franka-P-Rank dataset with the Robometer pipeline.
## Overview
- MIT-Franka-Prank is a robotics dataset with quality-labeled trajectories for manipulation tasks
- The dataset contains pre-recorded MP4 videos with metadata describing task instructions and quality labels
- Quality labels include: **successful**, **suboptimal**, and **failure**
## Dataset Structure
The dataset is organized as follows:
```
<dataset_path>/
20251210/
episode_000_foldtowel_suboptimal.mp4
episode_000_movebanana_success.mp4
episode_000_movepebble_suboptimal.mp4
...
foldtowel_metadata.json
pickandplace_metadata.json
```
### Tasks Included
1. **foldtowel**: Fold the towel in half
2. **movebanana**: Pick up the banana and place it on the blue plate
3. **movepebble**: Move some pebbles from the blue bowl to the green plate using the scoop
### Metadata Format
Each metadata JSON file contains an array of episodes:
```json
[
{
"episode_idx": 0,
"task_name": "foldtowel",
"filename": "episode_000_foldtowel_suboptimal.mp4",
"run_dir": "20251210_192953",
"instruction": "fold the towel in half",
"success": "suboptimal",
"trajectory_length": 1351
},
...
]
```
## Configuration
Configuration file: `dataset_upload/configs/data_gen_configs/mit_robot_prank.yaml`
```yaml
dataset:
dataset_path: ~/projects/robometer/datasets/20251210-mit-robot-prank
dataset_name: mit_franka_p-rank_rfm
output:
output_dir: ./robometer_dataset/mit_franka_p-rank_rfm
max_trajectories: -1 # -1 for all trajectories
max_frames: 64
use_video: true
fps: 10
shortest_edge_size: 240
center_crop: false
num_workers: 4
hub:
push_to_hub: true
hub_repo_id: mit_franka_p-rank_rfm
```
## Usage
### Convert Dataset to HuggingFace Format
```bash
uv run python -m dataset_upload.generate_hf_dataset --config_path=dataset_upload/configs/data_gen_configs/mit_franka_prank.yaml
```
This will:
- Read all metadata JSON files from the dataset directory
- Load the corresponding MP4 videos
- Process and resample videos to the specified frame count and FPS
- Generate language embeddings for task instructions
- Create a HuggingFace dataset with proper quality labels
- Optionally push to HuggingFace Hub
### Quality Label Mapping
The loader automatically normalizes quality labels:
- `"success"` → `"successful"`
- `"fail"` → `"failure"`
- `"suboptimal"` → `"suboptimal"` (unchanged)
## Output Format
The generated dataset will have the following schema:
- `id`: Unique trajectory identifier
- `task`: Task instruction text
- `lang_vector`: Language embedding of the task
- `data_source`: Dataset name
- `frames`: Path to video file (or sequence of images)
- `is_robot`: Boolean (True for this dataset)
- `quality_label`: One of "successful", "suboptimal", "failure"
- `preference_group_id`: None for this dataset
- `preference_rank`: None for this dataset
## Notes
- Videos are already in MP4 format, so the loader reads them directly using OpenCV
- The loader supports parallel processing with configurable worker count
- Language embeddings are cached to avoid redundant computations
- Output videos maintain the same content but are resampled to the specified frame count and FPS
## Troubleshooting
### Video File Not Found
- Ensure the dataset path points to the parent directory containing the date-stamped subdirectory
- Check that video files exist and match the filenames in metadata JSON files
### Missing Metadata Files
- The loader looks for files ending with `_metadata.json` in the video directory
- Ensure at least one metadata file exists (e.g., `foldtowel_metadata.json`, `pickandplace_metadata.json`)
### OpenCV Issues
- If you encounter video codec issues, ensure OpenCV is properly installed with video support
- Try: `pip install opencv-python-headless` or `pip install opencv-python`
## Example Output
After processing, you'll have a directory structure like:
```
robometer_dataset/mit_franka_p-rank_rfm/
mit_franka_p-rank_rfm/
shard_0000/
episode_000000/
foldtowel.mp4
episode_000001/
foldtowel.mp4
...
```
And a HuggingFace Dataset with all trajectory metadata.
|