Skip to content

wheelbot/dataset

Folders and files

NameName
Last commit message
Last commit date

Latest commit

Β 

History

72 Commits
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 

Repository files navigation

The Mini 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. The dataset is published on Zenodo.

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={Henrik Hose and Paul Brunzema and Devdutt Subhasish and Sebastian Trimpe},
    year={2026},
    url={https://arxiv.org/abs/2601.11394}, 
}

πŸ€– 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

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.

About

The Mini Wheelbot Dataset

Resources

License

Stars

Watchers

Forks

Packages

No packages published

Contributors 2

  •  
  •  

Languages