[CI] Fix torchcodec VideoEncoder uint8 requirement and docs build

- Convert video frames to uint8 in _write_video before passing to
  VideoEncoder, which requires uint8 input.
- Update MLFlow test_log_video to use torchcodec VideoDecoder for
  read-back verification when available.
- Install torchcodec in docs build environment to fix sphinx gallery
  failures when torchvision >= 0.22.

Made-with: Cursor

Author: vmoens

.github/workflows/docs.yml

@@ -44,7 +44,8 @@ jobs:
         # Install system dependencies (Ubuntu-based PyTorch image)
         apt-get update
         apt-get install -y libglfw3 libglfw3-dev libgl1-mesa-glx libgl1-mesa-dev libegl1-mesa-dev \
-          freeglut3-dev libglu1-mesa libegl1 mesa-utils xvfb wget git cmake
+          freeglut3-dev libglu1-mesa libegl1 mesa-utils xvfb wget git cmake ffmpeg \
+          libavcodec-dev libavformat-dev libavutil-dev libswscale-dev libswresample-dev
         
         # 2. upgrade pip, ninja and packaging
         python -m pip install --upgrade pip
@@ -65,10 +66,12 @@ jobs:
           # Use stable PyTorch for releases (tags and release/* branches)
           python -m pip install torch torchvision --quiet --root-user-action=ignore
           python -m pip install tensordict --quiet --root-user-action=ignore
+          python -m pip install torchcodec --quiet --root-user-action=ignore
         else
           # Use nightly PyTorch for main, nightly, and PRs
           python -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cpu -U --quiet --root-user-action=ignore
           python -m pip install git+https://github.com/pytorch/tensordict.git --quiet --root-user-action=ignore
+          BUILD_AGAINST_ALL_FFMPEG_FROM_S3=1 python -m pip install --no-build-isolation git+https://github.com/meta-pytorch/torchcodec.git --quiet --root-user-action=ignore
         fi
         
         # 6. Install doc requirements BEFORE TorchRL to ensure gymnasium is available

test/test_loggers.py

@@ -404,9 +404,18 @@ def test_log_video(self, steps, mlflow_fixture):
             videos_dir = client.download_artifacts(run_id, "videos", artifacts_dir)
             for i, video_name in enumerate(os.listdir(videos_dir)):
                 video_path = os.path.join(videos_dir, video_name)
-                loaded_video, _, _ = torchvision.io.read_video(
-                    video_path, pts_unit="sec", output_format="TCHW"
-                )
+                if _has_torchcodec:
+                    from torchcodec.decoders import VideoDecoder
+
+                    loaded_video = (
+                        VideoDecoder(video_path)
+                        .get_frames_in_range(start=0, stop=128)
+                        .data
+                    )
+                else:
+                    loaded_video, _, _ = torchvision.io.read_video(
+                        video_path, pts_unit="sec", output_format="TCHW"
+                    )
                 if steps:
                     assert torch.allclose(loaded_video.int(), videos[i].int(), rtol=0.1)
                 else:

torchrl/record/loggers/common.py

@@ -57,8 +57,8 @@ def _write_video(filename, video_array, **kwargs):  # noqa: F811
     preset = options.pop("preset", None)
     pixel_format = options.pop("pixel_format", None)
 
-    # VideoEncoder expects (N, C, H, W); callers pass (T, H, W, C)
-    video_array = video_array.permute(0, 3, 1, 2).contiguous()
+    # VideoEncoder expects (N, C, H, W) uint8; callers pass (T, H, W, C)
+    video_array = video_array.permute(0, 3, 1, 2).contiguous().to(torch.uint8)
 
     to_file_kwargs = {}
     if video_codec is not None: