DLR-RM / DLR-RM/stable-baselines3
[Bug]: Video upload to wandb broken since 2.4.0
- Dominant language
- Python
- Stars
- 13.8k
- Forks
- 2.2k
- Avg merge
- 1h 35m
- Merged PRs (30d)
- 2
Description
### 🐛 Bug
Using stable_baselines3 2.3.2 in Python 3.11 the provided unit test can upload videos to WANDB successfully. However, using 2.4 it fails.
### To Reproduce
```python
import unittest
import time
import os
import gymnasium as gym
from stable_baselines3.common.vec_env import VecVideoRecorder, DummyVecEnv
import wandb
from wandb import Api
from wandb.integration.sb3 import WandbCallback
from stable_baselines3 import PPO
class TestWandbVideoUpload(unittest.TestCase):
def test_video_upload(self):
env_id = "CartPole-v1"
video_folder = "videos"
video_length = 100
vec_env = DummyVecEnv([lambda: gym.make(env_id, render_mode="rgb_array")])
obs = vec_env.reset()
run = wandb.init(
project="test",
sync_tensorboard=True, # Automatically upload SB3's TensorBoard metrics
monitor_gym=True, # Automatically upload agent playing videos
# save_code=True, # Optional
)
# Record the video starting at the first step
vec_env = VecVideoRecorder(
vec_env,
video_folder,
record_video_trigger=lambda x: x == 0,
video_length=video_length,
name_prefix=f"agent-{env_id}"
)
vec_env.reset()
model = PPO("MlpPolicy", vec_env, verbose=1, tensorboard_log=f"runs/{run.id}")
model.learn(
total_timesteps=5000,
callback=WandbCallback(
model_save_path=f"tmp/models/{run.id}",
verbose=2,
),
)
run.finish()
# Give some time for the upload (adjust depending on connection speed)
time.sleep(30)
# Use the wandb API to check the run
api = Api()
# If you're logged into a different W&B account or using an organization, adjust 'entity' accordingly
run_path = f"{run.entity}/{run.project}/{run.id}"
run_api = api.run(run_path)
# Retrieve a list of all files in the run
files = run_api.files()
file_names = [f.name for f in files]
# Check if a video file is present
video_files = [name for name in file_names if name.endswith('.mp4')]
self.assertTrue(len(video_files) > 0, "The video was not uploaded to wandb.")
# Optional: Print the uploaded video files
print("Uploaded video files:", video_files)
# Clean up
vec_env.close()
wandb.finish()
if __name__ == '__main__':
unittest.main()
```
### Relevant log output / Error message
_No response_
### System Info
- OS: Linux-5.15.0-124-generic-x86_64-with-glibc2.35 # 134-Ubuntu SMP Fri Sep 27 20:20:17 UTC 2024
- Python: 3.11.0rc1
- Stable-Baselines3: 2.4.0
- PyTorch: 2.5.1+cu124
- GPU Enabled: False
- Numpy: 1.26.4
- Cloudpickle: 3.1.0
- Gymnasium: 0.29.1
### Checklist
- [X] My issue does not relate to a custom gym environment. (Use the custom gym env template instead)
- [X] I have checked that there is no similar [issue](https://github.com/DLR-RM/stable-baselines3/issues) in the repo
- [X] I have read the [documentation](https://stable-baselines3.readthedocs.io/en/master/)
- [X] I have provided a [minimal and working](https://github.com/DLR-RM/stable-baselines3/issues/982#issuecomment-1197044014) example to reproduce the bug
- [X] I've used the [markdown code blocks](https://help.github.com/en/articles/creating-and-highlighting-code-blocks) for both code and stack traces.
Contributor guide
Research direction
Start with the provided TestWandbVideoUpload reproduction and the WandbCallback integration, comparing behavior between Stable-Baselines3 2.3.2 and 2.4.0. Check the monitor_gym video-upload path and confirm the fix by verifying that the run contains an uploaded .mp4 file.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100