From d39e7f0bc4d5d4dcbe2aa6872e312f2fad5463a8 Mon Sep 17 00:00:00 2001 From: nemo Date: Tue, 14 Oct 2025 15:36:06 +0200 Subject: [PATCH] Log adapter config to WandB --- src/lerobot/rl/wandb_utils.py | 27 ++++++++++++++++++++++++++- 1 file changed, 26 insertions(+), 1 deletion(-) diff --git a/src/lerobot/rl/wandb_utils.py b/src/lerobot/rl/wandb_utils.py index 01cef9487..8c92e7145 100644 --- a/src/lerobot/rl/wandb_utils.py +++ b/src/lerobot/rl/wandb_utils.py @@ -112,7 +112,32 @@ class WandBLogger: artifact_name = f"{self._group}-{step_id}" artifact_name = get_safe_wandb_artifact_name(artifact_name) artifact = self._wandb.Artifact(artifact_name, type="model") - artifact.add_file(checkpoint_dir / PRETRAINED_MODEL_DIR / SAFETENSORS_SINGLE_FILE) + pretrained_model_dir = checkpoint_dir / PRETRAINED_MODEL_DIR + + # Check if this is a PEFT model (has adapter files instead of model.safetensors) + adapter_model_file = pretrained_model_dir / "adapter_model.safetensors" + standard_model_file = pretrained_model_dir / SAFETENSORS_SINGLE_FILE + + if adapter_model_file.exists(): + # PEFT model: add adapter files and configs + artifact.add_file(adapter_model_file) + adapter_config_file = pretrained_model_dir / "adapter_config.json" + if adapter_config_file.exists(): + artifact.add_file(adapter_config_file) + # Also add the policy config which is needed for loading + config_file = pretrained_model_dir / "config.json" + if config_file.exists(): + artifact.add_file(config_file) + elif standard_model_file.exists(): + # Standard model: add the single safetensors file + artifact.add_file(standard_model_file) + else: + logging.warning( + f"No {SAFETENSORS_SINGLE_FILE} or adapter_model.safetensors found in {pretrained_model_dir}. " + "Skipping model artifact upload to WandB." + ) + return + self._wandb.log_artifact(artifact) def log_dict(