Simplify configs (#550)

Co-authored-by: Remi <remi.cadene@huggingface.co>
Co-authored-by: HUANG TZU-CHUN <137322177+tc-huang@users.noreply.github.com>

lerobot/common/logger.py

@@ -21,32 +21,39 @@
 import logging
 import os
 import re
+from dataclasses import asdict
 from glob import glob
 from pathlib import Path
 
+import draccus
 import torch
 from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE
-from omegaconf import DictConfig, OmegaConf
 from termcolor import colored
 from torch.optim import Optimizer
 from torch.optim.lr_scheduler import LRScheduler
 
-from lerobot.common.policies.policy_protocol import Policy
-from lerobot.common.utils.utils import get_global_random_state, set_global_random_state
+from lerobot.common.policies.pretrained import PreTrainedPolicy
+from lerobot.common.utils.utils import get_global_random_state
+from lerobot.configs.train import TrainPipelineConfig
+from lerobot.configs.types import FeatureType, NormalizationMode
+
+PRETRAINED_MODEL = "pretrained_model"
+TRAINING_STATE = "training_state.pth"
 
 
 def log_output_dir(out_dir):
     logging.info(colored("Output dir:", "yellow", attrs=["bold"]) + f" {out_dir}")
 
 
-def cfg_to_group(cfg: DictConfig, return_list: bool = False) -> list[str] | str:
+def cfg_to_group(cfg: TrainPipelineConfig, return_list: bool = False) -> list[str] | str:
     """Return a group name for logging. Optionally returns group name as list."""
     lst = [
-        f"policy:{cfg.policy.name}",
-        f"dataset:{cfg.dataset_repo_id}",
-        f"env:{cfg.env.name}",
+        f"policy:{cfg.policy.type}",
+        f"dataset:{cfg.dataset.repo_id}",
         f"seed:{cfg.seed}",
     ]
+    if cfg.env is not None:
+        lst.append(f"env:{cfg.env.type}")
     return lst if return_list else "-".join(lst)
 
 
@@ -68,7 +75,6 @@ class Logger:
     The logger creates the following directory structure:
 
     provided_log_dir
-    ├── .hydra  # hydra's configuration cache
     ├── checkpoints
     │   ├── specific_checkpoint_name
     │   │   ├── pretrained_model  # Hugging Face pretrained model directory
@@ -80,28 +86,21 @@ class Logger:
     │   └── last  # a softlink to the last logged checkpoint
     """
 
-    pretrained_model_dir_name = "pretrained_model"
-    training_state_file_name = "training_state.pth"
+    pretrained_model_dir_name = PRETRAINED_MODEL
+    training_state_file_name = TRAINING_STATE
 
-    def __init__(self, cfg: DictConfig, log_dir: str, wandb_job_name: str | None = None):
-        """
-        Args:
-            log_dir: The directory to save all logs and training outputs to.
-            job_name: The WandB job name.
-        """
+    def __init__(self, cfg: TrainPipelineConfig):
         self._cfg = cfg
-        self.log_dir = Path(log_dir)
+        self.log_dir = cfg.output_dir
         self.log_dir.mkdir(parents=True, exist_ok=True)
-        self.checkpoints_dir = self.get_checkpoints_dir(log_dir)
-        self.last_checkpoint_dir = self.get_last_checkpoint_dir(log_dir)
-        self.last_pretrained_model_dir = self.get_last_pretrained_model_dir(log_dir)
+        self.job_name = cfg.job_name
+        self.checkpoints_dir = self.get_checkpoints_dir(self.log_dir)
+        self.last_checkpoint_dir = self.get_last_checkpoint_dir(self.log_dir)
+        self.last_pretrained_model_dir = self.get_last_pretrained_model_dir(self.log_dir)
 
         # Set up WandB.
         self._group = cfg_to_group(cfg)
-        project = cfg.get("wandb", {}).get("project")
-        entity = cfg.get("wandb", {}).get("entity")
-        enable_wandb = cfg.get("wandb", {}).get("enable", False)
-        run_offline = not enable_wandb or not project
+        run_offline = not cfg.wandb.enable or not cfg.wandb.project
         if run_offline:
             logging.info(colored("Logs will be saved locally.", "yellow", attrs=["bold"]))
             self._wandb = None
@@ -115,13 +114,13 @@ class Logger:
 
             wandb.init(
                 id=wandb_run_id,
-                project=project,
-                entity=entity,
-                name=wandb_job_name,
-                notes=cfg.get("wandb", {}).get("notes"),
+                project=cfg.wandb.project,
+                entity=cfg.wandb.entity,
+                name=self.job_name,
+                notes=cfg.wandb.notes,
                 tags=cfg_to_group(cfg, return_list=True),
-                dir=log_dir,
-                config=OmegaConf.to_container(cfg, resolve=True),
+                dir=self.log_dir,
+                config=asdict(self._cfg),
                 # TODO(rcadene): try set to True
                 save_code=False,
                 # TODO(rcadene): split train and eval, and run async eval with job_type="eval"
@@ -150,17 +149,19 @@ class Logger:
         """
         return cls.get_last_checkpoint_dir(log_dir) / cls.pretrained_model_dir_name
 
-    def save_model(self, save_dir: Path, policy: Policy, wandb_artifact_name: str | None = None):
+    def save_model(self, save_dir: Path, policy: PreTrainedPolicy, wandb_artifact_name: str | None = None):
         """Save the weights of the Policy model using PyTorchModelHubMixin.
 
         The weights are saved in a folder called "pretrained_model" under the checkpoint directory.
 
         Optionally also upload the model to WandB.
         """
+
         self.checkpoints_dir.mkdir(parents=True, exist_ok=True)
+        register_features_types()
         policy.save_pretrained(save_dir)
-        # Also save the full Hydra config for the env configuration.
-        OmegaConf.save(self._cfg, save_dir / "config.yaml")
+        # Also save the full config for the env configuration.
+        self._cfg.save_pretrained(save_dir)
         if self._wandb and not self._cfg.wandb.disable_artifact:
             # note wandb artifact does not accept ":" or "/" in its name
             artifact = self._wandb.Artifact(wandb_artifact_name, type="model")
@@ -173,18 +174,18 @@ class Logger:
         self,
         save_dir: Path,
         train_step: int,
-        optimizer: Optimizer,
-        scheduler: LRScheduler | None,
+        optimizer: Optimizer | None = None,
+        scheduler: LRScheduler | None = None,
     ):
         """Checkpoint the global training_step, optimizer state, scheduler state, and random state.
 
         All of these are saved as "training_state.pth" under the checkpoint directory.
         """
-        training_state = {
-            "step": train_step,
-            "optimizer": optimizer.state_dict(),
-            **get_global_random_state(),
-        }
+        training_state = {}
+        training_state["step"] = train_step
+        training_state.update(get_global_random_state())
+        if optimizer is not None:
+            training_state["optimizer"] = optimizer.state_dict()
         if scheduler is not None:
             training_state["scheduler"] = scheduler.state_dict()
         torch.save(training_state, save_dir / self.training_state_file_name)
@@ -192,10 +193,10 @@ class Logger:
     def save_checkpoint(
         self,
         train_step: int,
-        policy: Policy,
-        optimizer: Optimizer,
-        scheduler: LRScheduler | None,
         identifier: str,
+        policy: PreTrainedPolicy,
+        optimizer: Optimizer | None = None,
+        scheduler: LRScheduler | None = None,
     ):
         """Checkpoint the model weights and the training state."""
         checkpoint_dir = self.checkpoints_dir / str(identifier)
@@ -208,26 +209,11 @@ class Logger:
             checkpoint_dir / self.pretrained_model_dir_name, policy, wandb_artifact_name=wandb_artifact_name
         )
         self.save_training_state(checkpoint_dir, train_step, optimizer, scheduler)
-        os.symlink(checkpoint_dir.absolute(), self.last_checkpoint_dir)
 
-    def load_last_training_state(self, optimizer: Optimizer, scheduler: LRScheduler | None) -> int:
-        """
-        Given the last checkpoint in the logging directory, load the optimizer state, scheduler state, and
-        random state, and return the global training step.
-        """
-        training_state = torch.load(self.last_checkpoint_dir / self.training_state_file_name)
-        optimizer.load_state_dict(training_state["optimizer"])
-        if scheduler is not None:
-            scheduler.load_state_dict(training_state["scheduler"])
-        elif "scheduler" in training_state:
-            raise ValueError(
-                "The checkpoint contains a scheduler state_dict, but no LRScheduler was provided."
-            )
-        # Small hack to get the expected keys: use `get_global_random_state`.
-        set_global_random_state({k: training_state[k] for k in get_global_random_state()})
-        return training_state["step"]
+        relative_target = checkpoint_dir.relative_to(self.last_checkpoint_dir.parent)
+        self.last_checkpoint_dir.symlink_to(relative_target)
 
-    def log_dict(self, d, step, mode="train"):
+    def log_dict(self, d: dict, step: int, mode: str = "train"):
         assert mode in {"train", "eval"}
         # TODO(alexander-soare): Add local text log.
         if self._wandb is not None:
@@ -242,5 +228,13 @@ class Logger:
     def log_video(self, video_path: str, step: int, mode: str = "train"):
         assert mode in {"train", "eval"}
         assert self._wandb is not None
-        wandb_video = self._wandb.Video(video_path, fps=self._cfg.fps, format="mp4")
+        wandb_video = self._wandb.Video(video_path, fps=self._cfg.env.fps, format="mp4")
         self._wandb.log({f"{mode}/video": wandb_video}, step=step)
+
+
+def register_features_types():
+    draccus.decode.register(FeatureType, lambda x: FeatureType[x])
+    draccus.encode.register(FeatureType, lambda x: x.name)
+
+    draccus.decode.register(NormalizationMode, lambda x: NormalizationMode[x])
+    draccus.encode.register(NormalizationMode, lambda x: x.name)