diff --git a/amp_rsl_rl/runners/amp_on_policy_runner.py b/amp_rsl_rl/runners/amp_on_policy_runner.py
index 8e88862..9b31996 100644
--- a/amp_rsl_rl/runners/amp_on_policy_runner.py
+++ b/amp_rsl_rl/runners/amp_on_policy_runner.py
@@ -270,7 +270,7 @@ def update_run_name_with_sequence(prefix: str) -> None:
                 self.writer = WandbSummaryWriter(
                     log_dir=self.log_dir, flush_secs=10, cfg=self.cfg
                 )
-                update_run_name_with_sequence(prefix=self.cfg["wandb_project"])
+                update_run_name_with_sequence(prefix=self.cfg["wandb_kwargs"]["project"])
 
                 self.writer.log_config(
                     self.env.cfg, self.cfg, self.alg_cfg, self.policy_cfg
diff --git a/amp_rsl_rl/utils/wandb_utils.py b/amp_rsl_rl/utils/wandb_utils.py
index dea7945..621c328 100644
--- a/amp_rsl_rl/utils/wandb_utils.py
+++ b/amp_rsl_rl/utils/wandb_utils.py
@@ -40,6 +40,11 @@ def __init__(self, log_dir: str, flush_secs: int, cfg: dict) -> None:
         # Add log directory to wandb
         wandb.config.update({"log_dir": log_dir})
 
+        self.name_map = {
+            "Train/mean_reward/time": "Train/mean_reward_time",
+            "Train/mean_episode_length/time": "Train/mean_episode_length_time",
+        }
+
         self.video_files = []
 
     # To save video files to wandb explicitly