huggingface / lerobot

šŸ¤— LeRobot: Making AI for Robotics more accessible with end-to-end learning
Apache License 2.0
7.62k stars 726 forks source link

Update lerobot visualizer to send videos directly to Rerun #476

Open jleibs opened 1 month ago

jleibs commented 1 month ago

What this does

Rather than using the lerobot loader to decode the video and sending individual images, this now sends the video files directly.

To support this I had to introduce a new video_backend called raw which skips the dataset decoding.

I introduced a new flag --decode-video to maintain the old behavior.

On my desktop this

How it was tested

New functionality:

python lerobot/scripts/visualize_dataset.py  --repo-id lerobot/aloha_static_coffee_new  --episode-index 0

Old functionality:

python lerobot/scripts/visualize_dataset.py  --repo-id lerobot/aloha_static_coffee_new  --episode-index 0 --decode-video

Also verified old image pathway still works on dataset without video:

python lerobot/scripts/visualize_dataset.py     --repo-id lerobot/xarm_lift_medium_image     --episode-index 0

Crude benchmark

This results in a 10x real speedup, a 100x processing-overhead speedup, and a 50x reduction in .rrd file size.

Old pathway:

āžœ time python lerobot/scripts/visualize_dataset.py --repo-id lerobot/aloha_static_coffee_new --episode-index 0 --decode-video --save 1 --output-dir tmp

real    0m40.573s
user    7m33.682s
sys 0m53.324s

āžœ  ls -lh tmp/lerobot_aloha_static_coffee_new_episode_0.rrd 
-rw-r--r-- 1 jleibs jleibs 2.4G Oct 17 17:29 tmp/lerobot_aloha_static_coffee_new_episode_0.rrd

New pathway:

āžœ time python lerobot/scripts/visualize_dataset.py --repo-id lerobot/aloha_static_coffee_new --episode-index 0 --save 1 --output-dir tmp

real    0m4.380s
user    0m4.481s
sys 0m2.427s

āžœ  ls -lh tmp/lerobot_aloha_static_coffee_new_episode_0.rrd 
-rw-r--r-- 1 jleibs jleibs 48M Oct 17 17:31 tmp/lerobot_aloha_static_coffee_new_episode_0.rrd