From e23bb6819a4a2bada90a1664dcb76f44907f5b03 Mon Sep 17 00:00:00 2001 From: Farbod Farshidian Date: Wed, 7 Feb 2024 10:53:31 -0500 Subject: [PATCH 1/4] Check for store_code_state to store git-diff --- rsl_rl/runners/on_policy_runner.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rsl_rl/runners/on_policy_runner.py b/rsl_rl/runners/on_policy_runner.py index 9e0a459..6ae0837 100644 --- a/rsl_rl/runners/on_policy_runner.py +++ b/rsl_rl/runners/on_policy_runner.py @@ -155,7 +155,7 @@ def learn(self, num_learning_iterations: int, init_at_random_ep_len: bool = Fals if it % self.save_interval == 0: self.save(os.path.join(self.log_dir, f"model_{it}.pt")) ep_infos.clear() - if it == start_iter: + if it == start_iter and self.cfg.get("store_code_state", True): # obtain all the diff files git_file_paths = store_code_state(self.log_dir, self.git_status_repos) # if possible store them to wandb From 0a23fff25e7263cf1d4cd29ff6fb71e1c6f35751 Mon Sep 17 00:00:00 2001 From: Farbod Farshidian Date: Wed, 7 Feb 2024 20:00:01 -0500 Subject: [PATCH 2/4] Add video recording to wandb logger --- rsl_rl/runners/on_policy_runner.py | 4 +++ rsl_rl/utils/wandb_utils.py | 41 +++++++++++++++++++++++++++--- 2 files changed, 42 insertions(+), 3 deletions(-) diff --git a/rsl_rl/runners/on_policy_runner.py b/rsl_rl/runners/on_policy_runner.py index 6ae0837..08f91b6 100644 --- a/rsl_rl/runners/on_policy_runner.py +++ b/rsl_rl/runners/on_policy_runner.py @@ -210,6 +210,10 @@ def log(self, locs: dict, width: int = 80, pad: int = 35): "Train/mean_episode_length/time", statistics.mean(locs["lenbuffer"]), self.tot_time ) + # Video recording for wandb + if self.logger_type == "wandb": + self.writer.update_video_files(log_name="Video", fps=30.0) + str = f" \033[1m Learning iteration {locs['it']}/{locs['tot_iter']} \033[0m " if len(locs["rewbuffer"]) > 0: diff --git a/rsl_rl/utils/wandb_utils.py b/rsl_rl/utils/wandb_utils.py index 2868ce9..aff69f3 100644 --- a/rsl_rl/utils/wandb_utils.py +++ b/rsl_rl/utils/wandb_utils.py @@ -4,6 +4,8 @@ from __future__ import annotations import os +import json +import pathlib from dataclasses import asdict from torch.utils.tensorboard import SummaryWriter @@ -27,14 +29,15 @@ def __init__(self, log_dir: str, flush_secs: int, cfg): try: entity = os.environ["WANDB_USERNAME"] except KeyError: - raise KeyError( - "Wandb username not found. Please run or add to ~/.bashrc: export WANDB_USERNAME=YOUR_USERNAME" - ) + entity = None + print("Wandb username not found. Using default.") wandb.init(project=project, entity=entity) # Change generated name to project-number format wandb.run.name = project + wandb.run.name.split("-")[-1] + with open(os.path.join(log_dir, "wandb_info.json"), "w") as f: + json.dump({"wandb_run_id": wandb.run.id, "wandb_run_name": wandb.run.name}, f) self.name_map = { "Train/mean_reward/time": "Train/mean_reward_time", @@ -45,12 +48,18 @@ def __init__(self, log_dir: str, flush_secs: int, cfg): wandb.log({"log_dir": run_name}) + # Video logging. + self.saved_video_files = {} + def store_config(self, env_cfg, runner_cfg, alg_cfg, policy_cfg): wandb.config.update({"runner_cfg": runner_cfg}) wandb.config.update({"policy_cfg": policy_cfg}) wandb.config.update({"alg_cfg": alg_cfg}) wandb.config.update({"env_cfg": asdict(env_cfg)}) + def get_config(self): + return wandb.config + def _map_path(self, path): if path in self.name_map: return self.name_map[path] @@ -67,6 +76,29 @@ def add_scalar(self, tag, scalar_value, global_step=None, walltime=None, new_sty ) wandb.log({self._map_path(tag): scalar_value}, step=global_step) + def update_video_files(self, log_name, fps: float): + # Check if there are new video files + log_dir = pathlib.Path(self.log_dir) + video_files = list(log_dir.rglob("*.mp4")) + for video_file in video_files: + file_size_kb = os.stat(str(video_file)).st_size / 1024 + # If it is new file + if str(video_file) not in self.saved_video_files: + self.saved_video_files[str(video_file)] = {"size": file_size_kb, "added": False, "count": 0} + else: + # Only upload if the file size is not changing anymore to avoid uploading non-ready video. + video_info = self.saved_video_files[str(video_file)] + if video_info["size"] == file_size_kb and file_size_kb > 100 and not video_info["added"]: + if video_info["count"] > 10: + print(f"[Wandb] Uploading {os.path.basename(str(video_file))}.") + self.add_video(str(video_file), fps=fps, log_name=log_name) + self.saved_video_files[str(video_file)]["added"] = True + else: + video_info["count"] += 1 + else: + self.saved_video_files[str(video_file)]["size"] = file_size_kb + video_info["count"] = 0 + def stop(self): wandb.finish() @@ -78,3 +110,6 @@ 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(self, video_path: str, fps: int = 4, log_name: str = "Video"): + wandb.log({log_name: wandb.Video(video_path, fps=fps)}) From 64c2d15b4177068613f0c0fc288907dc86203004 Mon Sep 17 00:00:00 2001 From: Farbod Farshidian Date: Tue, 27 Feb 2024 13:12:52 -0500 Subject: [PATCH 3/4] Simplify --- rsl_rl/runners/on_policy_runner.py | 2 +- rsl_rl/utils/wandb_utils.py | 13 +++++-------- 2 files changed, 6 insertions(+), 9 deletions(-) diff --git a/rsl_rl/runners/on_policy_runner.py b/rsl_rl/runners/on_policy_runner.py index 08f91b6..937b36d 100644 --- a/rsl_rl/runners/on_policy_runner.py +++ b/rsl_rl/runners/on_policy_runner.py @@ -212,7 +212,7 @@ def log(self, locs: dict, width: int = 80, pad: int = 35): # Video recording for wandb if self.logger_type == "wandb": - self.writer.update_video_files(log_name="Video", fps=30.0) + self.writer.update_video_files(log_name="Video", fps=30) str = f" \033[1m Learning iteration {locs['it']}/{locs['tot_iter']} \033[0m " diff --git a/rsl_rl/utils/wandb_utils.py b/rsl_rl/utils/wandb_utils.py index aff69f3..e14661e 100644 --- a/rsl_rl/utils/wandb_utils.py +++ b/rsl_rl/utils/wandb_utils.py @@ -19,7 +19,7 @@ class WandbSummaryWriter(SummaryWriter): """Summary writer for Weights and Biases.""" def __init__(self, log_dir: str, flush_secs: int, cfg): - super().__init__(log_dir, flush_secs) + super().__init__(log_dir=log_dir, flush_secs=flush_secs) try: project = cfg["wandb_project"] @@ -30,7 +30,7 @@ def __init__(self, log_dir: str, flush_secs: int, cfg): entity = os.environ["WANDB_USERNAME"] except KeyError: entity = None - print("Wandb username not found. Using default.") + print("`WANDB_USERNAME` is not found! The run will be sent to your username.") wandb.init(project=project, entity=entity) @@ -76,7 +76,7 @@ def add_scalar(self, tag, scalar_value, global_step=None, walltime=None, new_sty ) wandb.log({self._map_path(tag): scalar_value}, step=global_step) - def update_video_files(self, log_name, fps: float): + def update_video_files(self, log_name: str, fps: int): # Check if there are new video files log_dir = pathlib.Path(self.log_dir) video_files = list(log_dir.rglob("*.mp4")) @@ -88,10 +88,10 @@ def update_video_files(self, log_name, fps: float): else: # Only upload if the file size is not changing anymore to avoid uploading non-ready video. video_info = self.saved_video_files[str(video_file)] - if video_info["size"] == file_size_kb and file_size_kb > 100 and not video_info["added"]: + if video_info["added"] is False and video_info["size"] == file_size_kb and file_size_kb > 100: if video_info["count"] > 10: print(f"[Wandb] Uploading {os.path.basename(str(video_file))}.") - self.add_video(str(video_file), fps=fps, log_name=log_name) + wandb.log({log_name: wandb.Video(str(video_file), fps=fps)}) self.saved_video_files[str(video_file)]["added"] = True else: video_info["count"] += 1 @@ -110,6 +110,3 @@ 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(self, video_path: str, fps: int = 4, log_name: str = "Video"): - wandb.log({log_name: wandb.Video(video_path, fps=fps)}) From 49e0a5b9bc167471cafaae05473653180524c9b2 Mon Sep 17 00:00:00 2001 From: Farbod Farshidian Date: Tue, 27 Feb 2024 13:34:57 -0500 Subject: [PATCH 4/4] Add description --- rsl_rl/utils/wandb_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/rsl_rl/utils/wandb_utils.py b/rsl_rl/utils/wandb_utils.py index e14661e..8552d9c 100644 --- a/rsl_rl/utils/wandb_utils.py +++ b/rsl_rl/utils/wandb_utils.py @@ -30,7 +30,7 @@ def __init__(self, log_dir: str, flush_secs: int, cfg): entity = os.environ["WANDB_USERNAME"] except KeyError: entity = None - print("`WANDB_USERNAME` is not found! The run will be sent to your username.") + print("`WANDB_USERNAME` is not found! WandB will request your username in the interactive mode.") wandb.init(project=project, entity=entity) @@ -48,7 +48,7 @@ def __init__(self, log_dir: str, flush_secs: int, cfg): wandb.log({"log_dir": run_name}) - # Video logging. + # Video logging bookkeeper self.saved_video_files = {} def store_config(self, env_cfg, runner_cfg, alg_cfg, policy_cfg):