Feat/wandb experiment tracking

README.md

@@ -215,6 +215,47 @@ algo.max_iterations=1000             # 覆盖迭代次数
 uv run python scripts/train_offpolicy.py --cfg job
 ```
 
+### W&B 实验记录
+
+将 `training.logger=wandb` 打开后，训练脚本会自动把实验记录到 Weights & Biases，并在本地 run 目录中写出：
+
+- `run_config.json`：完整运行配置
+- `run_summary.json`：运行摘要、耗时、最终效果
+
+MuJoCo 路径下如果训练后自动生成了 `play_video.mp4`，也会一并上传到当前 W&B run。
+
+```bash
+# 基本用法
+uv run python scripts/train_rsl_rl.py task=go1_joystick training.logger=wandb
+
+# 指定共享 project / entity
+uv run python scripts/train_appo.py \
+  task=go1_joystick \
+  training.logger=wandb \
+  training.wandb_project=unilab-benchmark \
+  training.wandb_entity=my-team
+
+# 同一个共享 project 下，按 task 分组
+uv run python scripts/train_offpolicy.py \
+  algo=sac \
+  task=go2_joystick \
+  training.logger=wandb \
+  training.wandb_project=unilab-benchmark \
+  training.wandb_group=go2_joystick
+```
+
+常用字段：
+
+- `training.wandb_project`：W&B project 名称，适合多人共享同一个实验空间
+- `training.wandb_entity`：团队或账号名
+- `training.wandb_group`：同一类实验的分组，建议按 task 或 benchmark 批次组织
+- `training.wandb_name`：自定义 run 名称；不填时会自动生成，避免多人共用 project 时重名
+- `training.wandb_tags`：额外标签，便于筛选实验
+- `training.wandb_notes`：备注
+- `training.wandb_mode=offline`：离线记录，后续再 `wandb sync`
+
+默认情况下，自动记录的信息包括任务、算法、backend、设备、硬件信息、git 信息、完整配置、整体运行时间、最终摘要指标，以及最终回放视频。
+
 ---
 
 ## APPO（异步 PPO）

conf/appo/config.yaml

@@ -52,6 +52,14 @@ training:
   device: null
   collector_device: null
   logger: tensorboard
+  wandb_project: unilab
+  wandb_entity: null
+  wandb_group: null
+  wandb_job_type: null
+  wandb_name: null
+  wandb_tags: []
+  wandb_notes: null
+  wandb_mode: null
   sim_backend: mujoco
   log_dir: null
   play_only: false

conf/offpolicy/config.yaml

@@ -8,6 +8,14 @@ training:
   task_name: Go1JoystickFlatTerrain
   device: null
   logger: tensorboard
+  wandb_project: unilab
+  wandb_entity: null
+  wandb_group: null
+  wandb_job_type: null
+  wandb_name: null
+  wandb_tags: []
+  wandb_notes: null
+  wandb_mode: null
   sim_backend: mujoco
   num_gpus: 1
   use_amp: false

conf/ppo/config.yaml

@@ -60,6 +60,14 @@ training:
   task_name: Go1JoystickFlatTerrain
   device: null
   logger: tensorboard
+  wandb_project: unilab
+  wandb_entity: null
+  wandb_group: null
+  wandb_job_type: null
+  wandb_name: null
+  wandb_tags: []
+  wandb_notes: null
+  wandb_mode: null
   sim_backend: mujoco
   play_only: false
   no_play: false

conf/ppo/config_mlx.yaml

@@ -58,6 +58,14 @@ training:
   play_env_num: 16
   play_steps: 150
   logger: tensorboard
+  wandb_project: unilab
+  wandb_entity: null
+  wandb_group: null
+  wandb_job_type: null
+  wandb_name: null
+  wandb_tags: []
+  wandb_notes: null
+  wandb_mode: null
   fp16: false
   log_root: logs/mlx_rl_train
   seed: 1

pyproject.toml

@@ -27,6 +27,7 @@ dependencies = [
     "typing-extensions",
     "mlx ; sys_platform == 'darwin'",
     "lark>=1.3.1",
+    "wandb",
 ]
 
 [project.optional-dependencies]

scripts/train_appo.py

@@ -16,6 +16,8 @@
 ROOT_DIR = Path(__file__).parent.parent
 sys.path.append(str(ROOT_DIR))
 
+from unilab.utils.experiment_tracking import ExperimentTracker
+
 
 def ensure_registries():
     for pkg_name in (
@@ -35,7 +37,7 @@ def ensure_registries():
             pass
 
 
-def play_appo(cfg: DictConfig, rl_cfg: dict):
+def play_appo(cfg: DictConfig, rl_cfg: dict) -> str | None:
     """Play mode for APPO."""
     import mediapy as media
     import numpy as np
@@ -139,7 +141,7 @@ def play_appo(cfg: DictConfig, rl_cfg: dict):
 
     if not load_path or not os.path.exists(load_path):
         print(f"Could not find run to load. load_path={load_path}")
-        return
+        return None
 
     print(f"Loading model: {load_path}")
     checkpoint = torch.load(load_path, map_location=device, weights_only=True)
@@ -183,6 +185,7 @@ def play_appo(cfg: DictConfig, rl_cfg: dict):
     print(f"Saving video to {output_video} with mediapy...")
     media.write_video(str(output_video), frames, fps=int(1.0 / env.cfg.ctrl_dt))
     print("Done.")
+    return output_video
 
 
 @hydra.main(version_base="1.3", config_path="../conf/appo", config_name="config")
@@ -208,40 +211,71 @@ def main(cfg: DictConfig) -> None:
     else:
         log_dir = cfg.training.log_dir
 
-    if not cfg.training.play_only:
-        from unilab.algos.torch.appo.runner import APPORunner
-
-        collector_device = cfg.training.collector_device
-        if collector_device == "gpu":
-            collector_device = "mps" if torch.backends.mps.is_available() else "cuda"
+    collector_device = cfg.training.collector_device
+    if collector_device == "gpu":
+        collector_device = "mps" if torch.backends.mps.is_available() else "cuda"
 
-        runner_kwargs = {}
-        if cfg.training.replay_queue_size is not None:
-            runner_kwargs["replay_queue_size"] = cfg.training.replay_queue_size
+    learner_device = cfg.training.device or (
+        "cuda"
+        if torch.cuda.is_available()
+        else "mps"
+        if torch.backends.mps.is_available()
+        else "cpu"
+    )
 
-        runner = APPORunner(
-            env_name=cfg.training.task_name,
-            env_cfg_overrides=env_cfg_override,
-            rl_cfg=rl_cfg,
-            device=cfg.training.device,
+    tracker = None
+    if not cfg.training.play_only:
+        tracker = ExperimentTracker(
+            root_dir=ROOT_DIR,
+            log_dir=log_dir,
+            algo_name="appo",
+            task_name=cfg.training.task_name,
+            sim_backend=cfg.training.sim_backend,
+            training_cfg=cfg.training,
+            full_cfg=cfg,
+            device=learner_device,
             collector_device=collector_device,
-            num_envs=cfg.algo.num_envs,
-            steps_per_env=cfg.algo.steps_per_env,
-            **runner_kwargs,
         )
-
-        try:
-            runner.learn(
-                max_iterations=cfg.algo.max_iterations,
-                save_interval=cfg.algo.save_interval,
-                log_dir=log_dir,
-                logger_type=cfg.training.logger,
+        tracker.start()
+
+    try:
+        if not cfg.training.play_only:
+            from unilab.algos.torch.appo.runner import APPORunner
+
+            runner_kwargs = {}
+            if cfg.training.replay_queue_size is not None:
+                runner_kwargs["replay_queue_size"] = cfg.training.replay_queue_size
+
+            runner = APPORunner(
+                env_name=cfg.training.task_name,
+                env_cfg_overrides=env_cfg_override,
+                rl_cfg=rl_cfg,
+                device=cfg.training.device,
+                collector_device=collector_device,
+                num_envs=cfg.algo.num_envs,
+                steps_per_env=cfg.algo.steps_per_env,
+                **runner_kwargs,
             )
-        finally:
-            runner.close()
 
-    if cfg.training.play_only or not cfg.training.no_play:
-        play_appo(cfg, rl_cfg)
+            try:
+                runner.learn(
+                    max_iterations=cfg.algo.max_iterations,
+                    save_interval=cfg.algo.save_interval,
+                    log_dir=log_dir,
+                    logger_type=cfg.training.logger,
+                )
+                if tracker is not None:
+                    tracker.update_summary(getattr(runner, "last_run_summary", None))
+            finally:
+                runner.close()
+
+        if cfg.training.play_only or not cfg.training.no_play:
+            play_video_path = play_appo(cfg, rl_cfg)
+            if tracker is not None:
+                tracker.log_video(play_video_path)
+    finally:
+        if tracker is not None:
+            tracker.finish()
 
 
 if __name__ == "__main__":