Updates wandb and neptune logging

rsl_rl/runners/on_policy_runner.py

@@ -156,7 +156,12 @@ def learn(self, num_learning_iterations: int, init_at_random_ep_len: bool = Fals
                 self.save(os.path.join(self.log_dir, f"model_{it}.pt"))
             ep_infos.clear()
             if it == start_iter:
-                store_code_state(self.log_dir, self.git_status_repos)
+                # obtain all the diff files
+                git_file_paths = store_code_state(self.log_dir, self.git_status_repos)
+                # if possible store them to wandb
+                if self.logger_type in ["wandb", "neptune"] and git_file_paths:
+                    for path in git_file_paths:
+                        self.writer.save_file(path)
 
         self.save(os.path.join(self.log_dir, f"model_{self.current_learning_iteration}.pt"))
 

rsl_rl/utils/neptune_utils.py

@@ -8,14 +8,14 @@
 from torch.utils.tensorboard import SummaryWriter
 
 try:
-    import neptune.new as neptune
+    import neptune
 except ModuleNotFoundError:
     raise ModuleNotFoundError("neptune-client is required to log to Neptune.")
 
 
 class NeptuneLogger:
     def __init__(self, project, token):
-        self.run = neptune.init(project=project, api_token=token)
+        self.run = neptune.init_run(project=project, api_token=token)
 
     def store_config(self, env_cfg, runner_cfg, alg_cfg, policy_cfg):
         self.run["runner_cfg"] = runner_cfg
@@ -86,3 +86,7 @@ def log_config(self, env_cfg, runner_cfg, alg_cfg, policy_cfg):
 
     def save_model(self, model_path, iter):
         self.neptune_logger.run["model/saved_model_" + str(iter)].upload(model_path)
+
+    def save_file(self, path, iter=None):
+        name = path.rsplit("/", 1)[-1].split(".")[0]
+        self.neptune_logger.run["git_diff/" + name].upload(path)

rsl_rl/utils/utils.py

@@ -58,15 +58,23 @@ def unpad_trajectories(trajectories, masks):
     )
 
 
-def store_code_state(logdir, repositories):
+def store_code_state(logdir, repositories) -> list:
+    git_log_dir = os.path.join(logdir, "git")
+    os.makedirs(git_log_dir, exist_ok=True)
+    file_paths = []
     for repository_file_path in repositories:
         try:
             repo = git.Repo(repository_file_path, search_parent_directories=True)
         except git.InvalidGitRepositoryError:
             # skip if not a git repository
             continue
+        # get the name of the repository
         repo_name = pathlib.Path(repo.working_dir).name
         t = repo.head.commit.tree
+        diff_file_name = os.path.join(git_log_dir, f"{repo_name}.diff")
         content = f"--- git status ---\n{repo.git.status()} \n\n\n--- git diff ---\n{repo.git.diff(t)}"
-        with open(os.path.join(logdir, f"{repo_name}_git.diff"), "x") as f:
+        with open(diff_file_name, "x") as f:
             f.write(content)
+        # add the file path to the list of files to be uploaded
+        file_paths.append(diff_file_name)
+    return file_paths

rsl_rl/utils/wandb_utils.py

@@ -74,4 +74,7 @@ def log_config(self, env_cfg, runner_cfg, alg_cfg, policy_cfg):
         self.store_config(env_cfg, runner_cfg, alg_cfg, policy_cfg)
 
     def save_model(self, model_path, iter):
-        wandb.save(model_path)
+        wandb.save(model_path, base_path=os.path.dirname(model_path))
+
+    def save_file(self, path, iter=None):
+        wandb.save(path, base_path=os.path.dirname(path))