What this does

This enables torchvision's — still experimental — video_reader backend for faster video decoding.

• In order to use it, torchvision has to be built from source using these instructions. This changes a few things in the dev docker image in order to be able to follow these instructions.
• Adds video_backend in config and as a LeRobotDataset (and MultiLeRobotDataset) argument to select between pyav and video_reader (defaults to pyav as before).
• Refactor the video benchmark.
• Adds changes from Add capture_camera_feed #267

How it was tested

python lerobot/common/datasets/_video_benchmark/run_video_benchmark.py 

with

BENCHMARKS = {
    "backend": ["pyav", "video_reader"],
}

Quality metrics (avg_per_pixel_l2_error, avg_psnr, avg_ssim, avg_mse) are identical. Loading time is generally improved by a factor of ~1.5.

1_frame

backend

2_frames

backend

2_frames_4_space

backend

6_frames

backend

I also did a run to reproduce pretrained act on aloha transfer task with

python lerobot/scripts/train.py \
  hydra.job.name=act_aloha_sim_transfer_cube_human \
  hydra.run.dir=outputs/train/act_aloha_sim_transfer_cube_human \
  policy=act \
  policy.use_vae=true \
  env=aloha \
  env.task=AlohaTransferCube-v0 \
  dataset_repo_id=lerobot/aloha_sim_transfer_cube_human \
  training.eval_freq=10000 \
  training.log_freq=250 \
  training.offline_steps=100000 \
  training.save_checkpoint=true \
  training.save_freq=25000 \
  eval.n_episodes=50 \
  eval.batch_size=50 \
  wandb.enable=true \
  device=cuda \
  video_backend=video_reader

The full run is available on wandb here

How to checkout & try? (for the reviewer)

You first need to compile torchvision from source. Original instructions are available there but I'll recap here as some of it is not up-to-date:

• Download the latest nvidia-video-codec-sdk and extract the zipped file.
• Set these environment variables needed for building torchvision against the sdk:
  • TORCHVISION_INCLUDE to the location of the video codec headers(nvcuvid.h and cuviddec.h), which would be under the Interface directory.
  • Set TORCHVISION_LIBRARY environment variable to the location of the video codec library(libnvcuvid.so), which would be under Lib/linux/stubs/x86_64 directory.
  • Set CUDA_HOME environment variable to the cuda root directory.

You can do all these with this command: (assuming you have unziped the codec sdk in $HOME and that your cuda home is at /usr/local/cuda, which is generally where it's at)

export TORCHVISION_INCLUDE=$HOME/Video_Codec_SDK_12.2.72/Interface/ \
    && export TORCHVISION_LIBRARY=$HOME/Video_Codec_SDK_12.2.72/Lib/linux/stubs/x86_64/ \
    && export CUDA_HOME=/usr/local/cuda

• Install ffmpeg<4.3 from conda-forge channel

conda install -c conda-forge "ffmpeg<4.3"

• Git-clone torchvision and do the following change in your pyproject.toml

- torchvision = ">=0.17.1"
+ torchvision = { path = "../../path/to/vision" }

• Run the following to update your dependencies and build torchvision

poetry lock --no-cache --no-update && poetry install --sync --all-extras

If you don't use poetry, you can simply do this instead (add the extras you need)

pip install .

• Ensure that video_reader works with this: (It shouldn't raise any error)

python -c "import torchvision; torchvision.set_video_backend('video_reader');"

Try any run with the option video_backend=video_reader (will default to pyav if not specified), e.g.

python lerobot/scripts/train.py \
    policy=act \
    policy.dim_model=64 \
    env=aloha \
    wandb.enable=False \
    training.offline_steps=2 \
    training.online_steps=0 \
    eval.n_episodes=1 \
    eval.batch_size=1 \
    device=cpu \
    training.save_checkpoint=true \
    training.save_freq=2 \
    policy.n_action_steps=20 \
    policy.chunk_size=20 \
    training.batch_size=2 \
    hydra.run.dir=tests/outputs/act/ \
    video_backend=video_reader

This change is