Enable video streaming and URDF logging from Rerun 0.24

[pweids]

Enable Video and URDF support from Rerun v0.24

What this does

This PR incorporates some new features to be released in Rerun 0.24 (that are currently available in the alpha release). Namely two things:

1. The ability to animate and visualize URDFs or robots
2. The ability to stream encoded video instead of just individual images

To utilize this, I've created some helpers classes/functions to modularize the logging procedures and bring them into record.py and teleoperate.py.

The main class here is the RerunRobotLogger in visualization_utils.py. This class is highly configurable, and makes it easy to log observation and action data.

On initialization, it will:

• initialize rerun
• load, parse, and log the URDF file
• set up any video streams with the right encoding settings

After initialization, the simplest way to use this is log_all, which will

• log all of the normalized observation joint values (current behavior)
• log all of the normalized action joint values (current behavior)
• convert the encoder steps to radians and log the joint angles
• log the captured images as either video, jpeg, or raw (depending on configuration)
• synchronize the time so that all logged data has the same timestamp

Video tradeoffs
There is a tradeoff between using video, compressed, and raw images, which is why I've made that configurable by the user. The tradeoff is essentially that there is a small encoding penalty for video, which means the live streaming of the video is about 1/2 second behind the robot. On the plus side, the entire video interaction can be played back at normal speeds.
With compressed images, there is also an encoding penalty, but it's negligible. This reduces the images from several megabytes to under 100 kilobytes. It's generally a best practice to log compressed images for this reason.
In both cases of logging the compressed images and the raw frames, I've kept static=True because of the memory issue. The video files remain a manageable size because of inter-frame video compression, so that does not cause much of a memory impact.

URDF logging
I've included the URDF files for the SO-100 and SO-101 arms. However, it's reasonable if you don't want to include this in the repository and instead instruct users to download them from somewhere else. These URDFs are slightly edited versions from here: https://github.com/TheRobotStudio/SO-ARM100/tree/main/Simulation

The RerunRobotLogger will try to load these URDFs, but an optional URDF path can be passed as well. And of course, URDF logging can be disabled. Logging the URDFs and meshes themselves is about 30-40MB of memory that's only incurred on launch, so some people may want to avoid this.

Additional changes

• in teleoperate.py and record.py, I made a few changes that were mostly to quiet the linter which didn't like the missing None checks. Happy to back these out if desired to reduce the footprint of this PR.
• I added an option in the config for the image stream type: "video", "jpeg", "raw", or "none"
• The video logging needs to be cleaned up to close the stream, so I've added a cleanup routine that must be run after the loops finish

How it was tested

I tested this mainly by running lerobot.teleoperate and lerobot.record manually with the various settings.

How to checkout & try? (for the reviewer)

To run this, version 0.24.0a4 is required. To avoid creating a dependency on an alpha version, I recommend we wait until 0.24 is released (targeting next week) before merging this in, so I'm going to submit this as a draft to get some eyes on it.

To run it, simply install the alpha version:

pip install rerun-sdk==0.24.0a4

Demo

I've uploaded a brief screen cap of running teleop with the URDF and video file here: https://drive.google.com/file/d/185APpSQl3rp3spfRhHAFwv9buWh2Grjj/view?usp=share_link

Known issues

There are some playback quirks we're working out. Specifically, when scrubbing to a random time, the video will rapidly playback from the most recent keyframe (basically displaying all the intermediary frames until the desired one). This will be resolved by the time 0.24 gets officially released.
There may also be better default encoding settings for reducing the encoding latency any further.

Follow-up work

There's clear duplicated work in that images are stored to disk raw and then encoded after each episode. With Rerun's live encoding, we should be able to save the encoded video files directly. This would save a lot of post-episode encoding time and disk IO.

[pweids]

note, the pre-commit bandit is failing because it thinks there's a typo in the element ID of one of the URDF files:

error: `ba` should be `by`, `be`
  --> lerobot/urdf/so101_follower/assets/wrist_roll_pitch_so101_v2.part:6:41
  |
6 |     "elementId": "3c22c2c23cb0ce545b9df2ba",
  |                                         ^^
  |

I bypassed this locally by committing with --no-verify (after solving all other pre-commit issues), but it's failing here. Not sure how to please the checker here