Skip to content

cwoodhayes/PolyUMI

Repository files navigation

lead.video.v1.1.-.slides_small.mov

PolyUMI: Visual + Auditory + Tactile Manipulation Platform for Imitation Learning

Project website: https://cwoodhayes.github.io/projects/polyumi
Hardware guide: Google Doc

PolyUMI is a real-time data collection & control platform for robotic imitation learning, which unifies the following sensor modalities in a single end-effector:

  • touch (via a custom optical tactile-sensing finger, based off of PolyTouch) - 10fps 540x480 MJPEG video (MP4)
  • mechanical vibration (via a contact microphone fixed to the finger housing) - 16kHz PCM audio (WAV)
  • vision (via GoPro camera on wrist + finger camera peripheral vision) - 60fps 1920x1080 MJPEG video (MP4) + 10fps 540x480 MJPEG video
  • proprioception (via monocular inertial SLAM from GoPro + IMU in gripper, or robot joint encoders + FK in embodiments)

It combines the Universal Manipulation Interface (UMI) platform with a custom touch-sensing finger inspired by the PolyTouch tactile + audio sensor, with hardware, firmware, and software built from scratch for a modern robotics stack (ROS2 Kilted + Python 3.13 + Foxglove visualizer).

Dataflow Overview

Data flow through the PolyUMI system.

Software Components

General summary of software components in this repo.

Repo Structure

pi/               # RPi client: camera, audio, LED streaming + episode recording
ingest/           # PC-side CLI: fetch sessions from Pi, encode video
inference_server/ # GPU-side FastAPI server for diffusion-policy inference (see docs/franka-inference-bringup.md)
ros2_ws/
  src/
    polyumi_pi_msgs/   # Protobuf message definitions (camera frame, audio chunk)
    polyumi_ros2/      # ROS 2 nodes + Foxglove launch files
external/
  ORB_SLAM3_PolyUMI/   # Git submodule: PolyUMI's ORB_SLAM3 fork (monocular visual-inertial SLAM for the GoPro Hero 12)

Prerequisites

PC: Python 3.13, uv, ROS 2 Kilted, ffmpeg, protobuf-compiler

RPi: Raspberry Pi Zero 2W flashed with Raspberry Pi OS. See docs/pi-provisioning.md for detailed setup instructions for the pi.

Installation

PC

Clone the repo and initialize the ORB-SLAM3 submodule (used by the SLAM preprocessing step):

git clone git@github.com:cwoodhayes/PolyUMI.git
cd PolyUMI
git submodule update --init --recursive

Install ingest dependencies (includes the polyumi_pi package for shared data types):

uv sync --group dev

Build the ROS 2 workspace:

cd ros2_ws
rosdep install --from-paths src --ignore-src -r --rosdistro kilted
colcon build
source install/setup.bash
cd ..

Build the ORB-SLAM3 fork. First, install the system dependencies (Ubuntu; the fork's README.md has more detail):

sudo apt install libopencv-dev libeigen3-dev libboost-serialization-dev

Then run the build script (~15 minutes; also builds the nested Pangolin, DBoW2, g2o, and Sophus submodules in-tree):

cd external/ORB_SLAM3_PolyUMI
bash build.sh
cd ../..

The ingest CLI then finds the binaries automatically — no env vars needed. Set ORB_SLAM3_DIR and ORB_SLAM3_BIN_SUBDIR only if pointing at an out-of-tree build.

RPi

Follow the instructions in docs/pi-provisioning.md to set up the pi for both gripper and end-effector.

After the setup instructions for the gripper above, the polyumi-pi systemd service will run every time the Pi boots, enabling you to record right away by pressing the button on the audio HAT.

Recording on the Gripper

  1. Turn on the GoPro attached to the UMI (until it's turned on, the Pi will not let you record)
  2. Turn on the Pi using the small button on the side of the PiSugar battery unit (short press, then long hold until all 4 LEDs light up, then release)
  3. Follow the original UMI data collection instructions (Step 0) to scan the configuration QR code for the GoPro.
  4. Wait until the red indicator LED on the audio HAT lights up red, indicating PolyUMI is ready to record. This may take 30-40 seconds after startup; the pi takes a while to boot.
  5. Press the button on the audio HAT to start recording an episode; the LED will pulse and the GoPro will start recording; press the button again to stop recording. Do not press the GoPro's shutter button or otherwise interact with the GoPro after powering it on; the pi will handle starting/stopping the GoPro's recording for synchronization.

Ingest

Data ingestion scripts that fetch recorded sessions from the Pi and pre-process them for use in training. From the repo root:

uv sync --group dev
uv tool install --editable ingest
# Show all ingest commands:
pingest --help
# Fetch only the latest session from the Pi:
pingest fetch --host <pi_ssh_hostname> --latest
# Fetch all new sessions:
pingest fetch --host <pi_ssh_hostname>
# Plug in the GoPro's SD card before running the fetch commands above to automatically fetch the GoPro's footage for each session
# OR download gopro footage from the SD later for all sessions you've fetched:
pingest fetch-gopro --host <pi_ssh_hostname>

# PROCESSING SESSIONS ON DISK
# ingest all new scenes on disk into their pzarr form, skipping already-processed scenes:
pingest process-all
# run the preprocessing pipeline on a particular scene (time alignment, SLAM, etc)
pingest pp <scene_directory> 
# export a single session to MCAP for easy visualization in foxglove (use the foxglove config 
# in ingest/foxglove)
pingest export-mcap <scene_directory> <session_number>
# export a scene's EPISODE sessions to a diffusion-policy ReplayBuffer zarr:
pingest export-dp <scene_directory> --output <output.zarr>

The at-rest data format used during the preprocessing stage managed by pingest is a zarr-based format stored in scene.zarr in each scene directory, referred to in these docs as pzarr. See docs/data-format.md for details on the format & the rationale behind it. The GoPro camera calibration shared by SLAM and the ArUco gripper-width step is documented in docs/camera-calibration.md.

Streaming / Demos

Streaming Demo

Streams camera, audio, and GoPro wrist camera into Foxglove.

On the Pi:

polyumi-pi stream

On the PC:

ros2 launch polyumi_ros2 stream_demo.launch.xml

Open Foxglove, connect to ws://localhost:8765, and drag in ros2_ws/src/polyumi_ros2/foxglove/layouts/stream_demo.json.

The launch file accepts two arguments: pi_host (default 10.106.10.62) and video_device (default /dev/video2) for the GoPro capture device.

Running on a robot arm

policy_client_node drives a robot arm from a diffusion-policy inference server. The arm's control stack typically runs on its own machine and is reached over ROS2; how you bring that up, network the two machines, and configure DDS depends on your robot and lab. docs/franka-inference-bringup.md describes the general inference architecture and API contract, and docs/crb-fr3-inference.md is a worked example for one specific Franka FR3 setup that you can adapt.

Hardware Notes

The Raspberry Pi app in pi/ also depends upon a custom fork of OpenGoPro with some specific extra capabilities & one bugfix; without these the library cannot connect over BLE on Debian Trixie (and its downsteam RPi fork which we are using.)

PiSugar Battery

Battery status is accessible at http://<pi_ip>:8421 or via I2C:

sudo i2cdetect -y 1
sudo i2cget -y 0x57 0x2a   # battery percentage; 100% = 0x64, 50% = 0x32, etc.

Troubleshooting

_version.py missing on the Pi — run ./deploy.sh <pi_ssh_hostname> from the PC; this generates the file from the current git HEAD.

Audio not detected — confirm wm8960-soundcard appears in arecord -l. If the default RaspiAudio driver was previously installed, the Waveshare DKMS driver may need to be reinstalled after a kernel update.

Wi-Fi not listing on the Pi — run sudo modprobe brcmfmac, then retry nmcli dev wifi connect "your-network".

ZMQ frames dropping — check the cb_drops counter in the Pi logs. The audio streamer uses a 100-frame queue with drop-and-replace on overflow; the video streamer uses NOBLOCK sends with a high-watermark of 2. Persistent drops indicate the network link is the bottleneck.

protoc not found during polyumi_pi_msgs install — install protobuf-compiler (sudo apt install protobuf-compiler on the Pi, or via your system package manager on the PC).

PiSugar 3 battery board won't turn off - the pisugar-server service has probably crashed. SSH into the Pi and run sudo systemctl restart pisugar-server to fix it. You can also interactively communicate with the service with nc -U /tmp/pisugar-server.sock; then send help to show commands. This has worked for me in all my cases; if this doesn't work for you, check the PiSugar documentation for more ideas.

arecord -D hw:wm8960soundcard -r 48000 -f S16_LE -c 2 -d 5 test.wav

Development

Developing the polyumi-pi app

To deploy the latest code to the Pi, run from the PC:

./deploy.sh <pi_ssh_hostname>

Then on the Pi, install the Python environment and the polyumi-pi script:

cd ~/PolyUMI/pi
# venv should have already been created by cloud-init, but run this if you need to recreate it for any reason (e.g. to pull in new system packages like picamera2):
# uv venv --system-site-packages
uv sync --no-dev --extra pi
source .venv/bin/activate

The pi extra pulls in the Raspberry Pi hardware-only dependencies (lgpio, gpiozero, rpi-hardware-pwm) that the full app needs at runtime. They are gated behind this extra because lgpio can't build off-Pi (it needs swig), so a plain PC dev sync omits them. deploy.sh passes --extra pi automatically.

picamera2 must be installed via apt, not pip — the --system-site-packages flag above pulls it in from the system. (this apt install & others is handled by the cloud-init provisioning workflow).

Tip for development: add the deploy.sh invocation to .vscode/tasks.json as a build task so it runs with Ctrl+Shift+B.

Run polyumi-pi --help for a full list of commands.

About

Visual-Auditory-Tactile Manipulation Data Collector for Imitation Learning

Resources

Stars

5 stars

Watchers

0 watching

Forks

Releases

No releases published

Packages

 
 
 

Contributors