[Experimental] Optimizes the save_episode()

[masafumi-airoa]

What this does

This PR experimentally improves the performance of the save_episode() function.

The encode_episode_videos() previously saved images as .png images, asynchronously loaded them as PIL images, converted them to RGB images, and created a video from the PIL images with PyAV. This process was the biggest bottleneck, accounting for approximately 56% of the total runtime (for converting 11 rosbags in my environment).

The new pipeline stores the images as np.ndarray in memory and directly inputs them to PyAV when frame_staging = "memory". This pipeline shortcuts the above process and reduces about 50% of the total execution time.

• Execution time (main.py with debug_data)
  • Before: 382.902 [sec]
  • After: 195.010 [sec]

How it was tested

Visually checked the generated videos.

hand
https://github.com/user-attachments/assets/9261391e-d0dc-4aab-8e2f-2cbe6de86648

head
https://github.com/user-attachments/assets/5f495929-62ce-4bdd-bb6c-529696323f1c

How to checkout & try? (for the reviewer)

We can measure the execution time by running these commands in the PR's branch after downloading debug_data dataset from AWS,

uv sync --reinstall-package lerobot
uv add pyinstrument
uv run pyinstrument -r html -o profile.html -m hsr_data_converter.rosbag2lerobot.main --raw_dir {path to the debug_data directory} --out_dir ./datasets --fps 10 --robot_type hsr --conversion_type aggregate --separate_per_primitive false