wheelbot-dataset 0.1.2

The Mini Wheelbot Dataset: High-Fidelity Data for Robot Learning

Wheelbot Dataset

A large, high-quality dynamics dataset of the Mini Wheelbot. The dataset contains 1 kHz data of all onboard sensor readings, the estimated state, ground-truth pose measurements from a motion capture system, and a third-person view video of the experiment. We perform a variety of experiments using pseudo-random binary excitation signals (PRBS) as setpoints of a linear controller, an MPC for driving, and an RL policy that races along a track. Experiments are performed across multiple hardware instances and on different surfaces. With this dataset, we hope to encourage researchers to use the Mini Wheelbot to benchmark their learning-based control methods, even without access to the real hardware. We include two example implementations of how to use the dataset, i.e., for dynamics learning and state estimation. A detailed description of this dataset is available in the brief on arxiv.

If you find this dataset helpful, please cite the Mini Wheelbot paper and/or dataset directly:

@inproceedings{hose2025mini,
   title={The {Mini Wheelbot}: A Testbed for Learning-based Balancing, Flips, and Articulated Driving},
   booktitle={2025 IEEE International Conference on Robotics and Automation (ICRA)},
   publisher={IEEE},
   author={Hose, Henrik and Weisgerber, Jan and Trimpe, Sebastian},
   year={2025},
}

@article{hose2026dataset,
  title={The {Mini Wheelbot} Dataset: High-Fidelity Data for Robot Learning},
  author={Hose, Henrik and Brunzema, Paul, and Subhasish, Devdutt and Trimpe, Sebastian},
  year={2026},
}

🤖 Overview

The Wheelbot Dataset contains trajectory data from a self-balancing wheeled robot across various experimental conditions including:

• Pitch experiments: Random pitch angle tracking
• Roll experiments: Random roll angle tracking
• Velocity experiments: Forward/backward velocity tracking
• Combined experiments: Velocity + roll, velocity + pitch
• Yaw experiments: Random yaw control, circular trajectories, figure-eight patterns

All data is recorded at 1000 Hz and includes:

• IMU data (4x gyroscopes, 4x accelerometers)
• Robot state (yaw, roll, pitch angles and velocities)
• Wheel positions, velocities, and accelerations
• Motor commands
• Setpoints
• Vicon motion capture data (ground truth)
• Battery voltage

📦 Installation

From PyPI (when published)

pip install wheelbot-dataset

From Source

git clone https://github.com/wheelbot/dataset.git
cd dataset
pip install -e .

Optional Dependencies

# For development and examples
pip install wheelbot-dataset[dev]

# For all features (neural network training, etc.)
pip install wheelbot-dataset[all]

🚀 Quick Start

1. Download the Dataset

The dataset is automatically downloadable from Zenodo:

# Download dataset to ./data directory
python -m wheelbot_dataset download download

# Or download to a custom location
python -m wheelbot_dataset download download --output_dir=my_dataset

# Check if dataset exists
python -m wheelbot_dataset download check

From Python:

from wheelbot_dataset import download_dataset, check_dataset

# Download the dataset
download_dataset(output_dir="data")

# Check dataset status
check_dataset("data")

2. Load and Explore Data

from wheelbot_dataset import Dataset

# Load the dataset
ds = Dataset("data")

# Load a specific experiment group
pitch_group = ds.load_group("pitch")

# Access individual experiments
exp = pitch_group[0]

# View the data
print(exp.data.head())  # Pandas DataFrame
print(exp.meta)         # Metadata dictionary

3. Process and Filter Data

# Cut off standup and laydown phases
exp_cut = exp.cut_by_condition(
    start_condition=lambda df: df['/tau_DR_command/reaction_wheel'].abs() > 0,
    end_condition=lambda df: df['/tau_DR_command/reaction_wheel'].abs() == 0
).cut_time(start=2.0, end=2.0)

# Apply filtering
import scipy.signal as signal

def lowpass_filter(df):
    b, a = signal.butter(4, 2*50/1000, btype="low")
    df_filtered = df.copy()
    for col in df.columns:
        df_filtered[col] = signal.filtfilt(b, a, df[col])
    return df_filtered

exp_filtered = exp_cut.apply_filter(lowpass_filter).resample(dt=0.01)

4. Batch Processing

# Filter entire groups or datasets
group_filtered = pitch_group.map(
    lambda exp: exp
        .filter_by_metadata(experiment_status="success")
        .cut_time(start=2.0, end=2.0)
        .apply_filter(lowpass_filter)
        .resample(dt=0.01)
)

# Filter by metadata
successful_exps = ds.map(
    lambda exp: exp.filter_by_metadata(experiment_status="success")
)

📊 Visualization

Plot Time Series

from wheelbot_dataset import plot_timeseries

plot_fields = {
    "Angles": ["/q_yrp/yaw", "/q_yrp/roll", "/q_yrp/pitch"],
    "Angular Velocities": ["/dq_yrp/yaw_vel", "/dq_yrp/roll_vel", "/dq_yrp/pitch_vel"],
    "Motor Commands": ["/tau_DR_command/drive_wheel", "/tau_DR_command/reaction_wheel"],
}

plot_timeseries(
    experiments=exp_filtered,
    groups=plot_fields,
    pdf_path="output.pdf"
)

Plot Histograms

from wheelbot_dataset import plot_histograms

plot_histograms(exp_filtered, "histograms.pdf")
plot_histograms(pitch_group, "group_histograms.pdf")
plot_histograms(ds, "dataset_histograms.pdf")

🤖 Machine Learning Integration

Export to NumPy

# Single experiment
columns = ["time"] + list(exp.data.columns)
exp_numpy = exp.to_numpy(columns)

# Entire group
group_numpy = group.map(lambda exp: exp.to_numpy(columns))

Create Training Dataset

from wheelbot_dataset import to_prediction_dataset

fields_states = [
    "/q_yrp/roll", "/q_yrp/pitch",
    "/dq_yrp/yaw_vel", "/dq_yrp/roll_vel", "/dq_yrp/pitch_vel",
    "/dq_DR/drive_wheel", "/dq_DR/reaction_wheel",
]
fields_actions = [
    "/tau_DR_command/drive_wheel", 
    "/tau_DR_command/reaction_wheel",
]

# One-step prediction dataset
states, actions, nextstates, _ = to_prediction_dataset(
    ds_filtered,
    fields_states=fields_states,
    fields_actions=fields_actions,
    N_future=1
)

# Multi-step prediction dataset
states, actions, nextstates, _ = to_prediction_dataset(
    ds_filtered,
    fields_states=fields_states,
    fields_actions=fields_actions,
    N_future=100
)

print(f"States shape: {states.shape}")
print(f"Actions shape: {actions.shape}")
print(f"Next states shape: {nextstates.shape}")

📈 Dataset Statistics

Get comprehensive statistics about your dataset:

python -m wheelbot_dataset consolidate statistics --dataset_path=data

This generates a LaTeX table with:

• Number of trajectories per experiment type
• Total duration (with configurable cutoff)
• Number of crashes/failed experiments

Analyze Update Rates

Check the actual update rates of different sensor signals:

python -m wheelbot_dataset consolidate updaterates \
    --dataset_path=data \
    --group_name=pitch \
    --index=4

🗂️ Dataset Structure

data/
├── pitch/              # Pitch angle tracking experiments
│   ├── 0.csv
│   ├── 0.meta
│   ├── 0.mp4
│   ├── ...
├── roll/               # Roll angle tracking experiments
├── roll_max/           # Maximum roll angle experiments
├── velocity/           # Forward/backward velocity
├── velocity_pitch/     # Combined velocity and pitch
├── velocity_roll/      # Combined velocity and roll
├── yaw/                # Random yaw control
├── yaw_circle/         # Circular trajectories
└── yaw_figure_eight/   # Figure-eight trajectories

Each experiment includes:

• .csv - Time series data at 1000 Hz
• .meta - JSON metadata (robot ID, surface, experiment status)
• .mp4 - Video recording
• .log - Raw log file
• .pkl - Setpoint data
• .preview.pdf - Quick visualization
• .setpoints.pdf - Setpoint visualization

🎓 Examples

Comprehensive examples are available in the examples/ directory:

Complementary Filter (examples/estimator/)

• Complementary filter for orientation estimation
• Comparison with onboard estimates and Vicon ground truth
• Run with: python examples/estimator/complementary_filter.py

Dynamics Model Learning (examples/dynamics_model_learning/)

• Neural network-based dynamics learning with JAX/Equinox
• One-step and multi-step prediction models with curriculum learning
• Dataset export and model training scripts
• Model evaluation and comparison plots

Time Series Classification (examples/timeseries_classification/)

• Transformer-based classification for surface detection
• Robot identification from IMU data
• Control mode classification (autonomous vs. human)

Basic Usage (examples/basic_usage.py)

A comprehensive example demonstrating the main features of the package:

# Run the basic usage example
python examples/basic_usage.py

The example script shows:

• Loading datasets and experiment groups
• Cutting and filtering data
• Applying custom filters (e.g., lowpass)
• Resampling data
• Batch processing with map operations
• Plotting time series and histograms
• Exporting to NumPy arrays
• Creating training datasets for machine learning

Data Fields

IMU Data

• /gyro{0-3}/{x,y,z} - Angular velocities from 4 gyroscopes
• /accel{0-3}/{x,y,z} - Linear accelerations from 4 accelerometers

Robot State

• /q_yrp/{yaw,roll,pitch} - Euler angles (rad)
• /dq_yrp/{yaw_vel,roll_vel,pitch_vel} - Angular velocities (rad/s)
• /q_DR/{drive_wheel,reaction_wheel} - Wheel positions (rad)
• /dq_DR/{drive_wheel,reaction_wheel} - Wheel velocities (rad/s)
• /ddq_DR/{drive_wheel,reaction_wheel} - Wheel accelerations (rad/s²)

Commands

• /tau_DR_command/{drive_wheel,reaction_wheel} - Motor torque commands (N⋅m)

Setpoints

• /setpoint/{yaw,roll,pitch} - Target Euler angles
• /setpoint/{yaw_rate,roll_rate,pitch_rate} - Target angular velocities
• /setpoint/driving_wheel_angle - Target wheel angle
• /setpoint/driving_wheel_angular_velocity - Target wheel velocity

Motion Capture

• /vicon_position/{x,y,z} - Position from Vicon (m)
• /vicon_orientation_wxyz/{w,x,y,z} - Quaternion orientation

Other

• battery/voltage - Battery voltage (V)

🛠️ Recording New Data

For information on recording new trajectories, see README_RECORDING.md.