From 2e62d044fba40bf12ff9f640a1adee90b2d60aa5 Mon Sep 17 00:00:00 2001 From: vaishnavahari seenivasan Date: Thu, 17 Apr 2025 13:25:06 +0000 Subject: [PATCH 1/4] upload videos to wandb --- rsl_rl/runners/on_policy_runner.py | 4 ++++ rsl_rl/utils/wandb_utils.py | 19 ++++++++++++++++++- 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/rsl_rl/runners/on_policy_runner.py b/rsl_rl/runners/on_policy_runner.py index ecda87ae..f40710c7 100644 --- a/rsl_rl/runners/on_policy_runner.py +++ b/rsl_rl/runners/on_policy_runner.py @@ -350,6 +350,10 @@ def log(self, locs: dict, width: int = 80, pad: int = 35): str = f" \033[1m Learning iteration {locs['it']}/{locs['tot_iter']} \033[0m " + # Upload video to wandb + if self.logger_type == "wandb" and not self.disable_logs: + self.writer.add_video_files(self.log_dir, fps=self.cfg["video_fps"], step=locs["it"]) + if len(locs["rewbuffer"]) > 0: log_string = ( f"""{'#' * width}\n""" diff --git a/rsl_rl/utils/wandb_utils.py b/rsl_rl/utils/wandb_utils.py index 243e82d4..647a071d 100644 --- a/rsl_rl/utils/wandb_utils.py +++ b/rsl_rl/utils/wandb_utils.py @@ -5,7 +5,7 @@ from __future__ import annotations -import os +import os, pathlib, json from dataclasses import asdict from torch.utils.tensorboard import SummaryWriter @@ -44,6 +44,7 @@ def __init__(self, log_dir: str, flush_secs: int, cfg): "Train/mean_reward/time": "Train/mean_reward_time", "Train/mean_episode_length/time": "Train/mean_episode_length_time", } + self.video_files = [] def store_config(self, env_cfg, runner_cfg, alg_cfg, policy_cfg): wandb.config.update({"runner_cfg": runner_cfg}) @@ -76,6 +77,22 @@ def save_model(self, model_path, iter): def save_file(self, path, iter=None): wandb.save(path, base_path=os.path.dirname(path)) + def add_video_files(self, log_dir: str, step: int, fps: int = 30): + # Check if there are video files in the video directory + videos_dir = os.path.join(log_dir, "videos") + if os.path.exists(videos_dir): + # append the new video files to the existing list + for video_file in os.listdir(videos_dir): + if video_file.endswith(".mp4") and video_file not in self.video_files: + self.video_files.append(video_file) + # add the new video file to wandb only if video file is not updating + video_path = os.path.join(videos_dir, video_file) + wandb.log( + {"Video": wandb.Video(video_path, fps=fps, format="mp4")}, + step = step + ) + + """ Private methods. """ From e1a6e36875e4e15c250857883dd2f517d3943a4b Mon Sep 17 00:00:00 2001 From: svaichu Date: Thu, 17 Apr 2025 14:25:16 +0000 Subject: [PATCH 2/4] Improve video upload functionality by using configurable fps or defaulting to 30 --- rsl_rl/runners/on_policy_runner.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/rsl_rl/runners/on_policy_runner.py b/rsl_rl/runners/on_policy_runner.py index f40710c7..f5f38199 100644 --- a/rsl_rl/runners/on_policy_runner.py +++ b/rsl_rl/runners/on_policy_runner.py @@ -352,7 +352,12 @@ def log(self, locs: dict, width: int = 80, pad: int = 35): # Upload video to wandb if self.logger_type == "wandb" and not self.disable_logs: - self.writer.add_video_files(self.log_dir, fps=self.cfg["video_fps"], step=locs["it"]) + # use video_fps from cfg if available or default to 30 + if "video_fps" in self.cfg: + video_fps = self.cfg["video_fps"] + else: + video_fps = 30 + self.writer.add_video_files(self.log_dir, fps=video_fps, step=locs["it"]) if len(locs["rewbuffer"]) > 0: log_string = ( From 31e600002cb9ec8a108c247c583fde0b020b4c33 Mon Sep 17 00:00:00 2001 From: svaichu Date: Thu, 17 Apr 2025 16:52:01 +0000 Subject: [PATCH 3/4] consider videos only from train directory --- rsl_rl/utils/wandb_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rsl_rl/utils/wandb_utils.py b/rsl_rl/utils/wandb_utils.py index 647a071d..1a829dab 100644 --- a/rsl_rl/utils/wandb_utils.py +++ b/rsl_rl/utils/wandb_utils.py @@ -79,7 +79,7 @@ def save_file(self, path, iter=None): def add_video_files(self, log_dir: str, step: int, fps: int = 30): # Check if there are video files in the video directory - videos_dir = os.path.join(log_dir, "videos") + videos_dir = os.path.join(log_dir, "videos", "train") if os.path.exists(videos_dir): # append the new video files to the existing list for video_file in os.listdir(videos_dir): From 1a83d610a43bf2018bf8556bf68f67dfc612286d Mon Sep 17 00:00:00 2001 From: svaichu Date: Tue, 22 Apr 2025 06:54:43 +0200 Subject: [PATCH 4/4] look for all subdirectories of log_dir --- rsl_rl/utils/wandb_utils.py | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/rsl_rl/utils/wandb_utils.py b/rsl_rl/utils/wandb_utils.py index 1a829dab..e4eec795 100644 --- a/rsl_rl/utils/wandb_utils.py +++ b/rsl_rl/utils/wandb_utils.py @@ -79,18 +79,18 @@ def save_file(self, path, iter=None): def add_video_files(self, log_dir: str, step: int, fps: int = 30): # Check if there are video files in the video directory - videos_dir = os.path.join(log_dir, "videos", "train") - if os.path.exists(videos_dir): + if os.path.exists(log_dir): # append the new video files to the existing list - for video_file in os.listdir(videos_dir): - if video_file.endswith(".mp4") and video_file not in self.video_files: - self.video_files.append(video_file) - # add the new video file to wandb only if video file is not updating - video_path = os.path.join(videos_dir, video_file) - wandb.log( - {"Video": wandb.Video(video_path, fps=fps, format="mp4")}, - step = step - ) + for root, dirs, files in os.walk(log_dir): + for video_file in files: + if video_file.endswith(".mp4") and video_file not in self.video_files: + self.video_files.append(video_file) + # add the new video file to wandb only if video file is not updating + video_path = os.path.join(root, video_file) + wandb.log( + {"Video": wandb.Video(video_path, fps=fps, format="mp4")}, + step = step + ) """