ydy0615/lerobot-clone
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-05-31 10:51:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

Log adapter config to WandB

Browse Source
This commit is contained in:
nemo
2025-10-14 15:36:06 +02:00
parent 20d4b5d347
commit d39e7f0bc4
1 changed files with 26 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

27
src/lerobot/rl/wandb_utils.py
View File

@@ -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(