Release cleanup (#132)

Co-authored-by: Kashif Rasul <kashif.rasul@gmail.com>
Co-authored-by: Alexander Soare <alexander.soare159@gmail.com>
Co-authored-by: Adil Zouitine <adilzouitinegm@gmail.com>
Co-authored-by: Cadene <re.cadene@gmail.com>

lerobot/scripts/eval.py

@@ -583,17 +583,18 @@ if __name__ == "__main__":
             pretrained_policy_path = Path(
                 snapshot_download(args.pretrained_policy_name_or_path, revision=args.revision)
             )
-        except HFValidationError:
-            logging.warning(
-                "The provided pretrained_policy_name_or_path is not a valid Hugging Face Hub repo ID. "
-                "Treating it as a local directory."
-            )
-        except RepositoryNotFoundError:
-            logging.warning(
-                "The provided pretrained_policy_name_or_path was not found on the Hugging Face Hub. Treating "
-                "it as a local directory."
-            )
-        pretrained_policy_path = Path(args.pretrained_policy_name_or_path)
+        except (HFValidationError, RepositoryNotFoundError) as e:
+            if isinstance(e, HFValidationError):
+                error_message = (
+                    "The provided pretrained_policy_name_or_path is not a valid Hugging Face Hub repo ID."
+                )
+            else:
+                error_message = (
+                    "The provided pretrained_policy_name_or_path was not found on the Hugging Face Hub."
+                )
+
+            logging.warning(f"{error_message} Treating it as a local directory.")
+            pretrained_policy_path = Path(args.pretrained_policy_name_or_path)
         if not pretrained_policy_path.is_dir() or not pretrained_policy_path.exists():
             raise ValueError(
                 "The provided pretrained_policy_name_or_path is not a valid/existing Hugging Face Hub "

lerobot/scripts/push_dataset_to_hub.py

@@ -60,7 +60,7 @@ import torch
 from huggingface_hub import HfApi
 from safetensors.torch import save_file
 
-from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.common.datasets.lerobot_dataset import CODEBASE_VERSION, LeRobotDataset
 from lerobot.common.datasets.push_dataset_to_hub._download_raw import download_raw
 from lerobot.common.datasets.push_dataset_to_hub.compute_stats import compute_stats
 from lerobot.common.datasets.utils import flatten_dict
@@ -252,7 +252,7 @@ def main():
     parser.add_argument(
         "--revision",
         type=str,
-        default="v1.2",
+        default=CODEBASE_VERSION,
         help="Codebase version used to generate the dataset.",
     )
     parser.add_argument(

lerobot/scripts/train.py

@@ -8,7 +8,6 @@ import hydra
 import torch
 from datasets import concatenate_datasets
 from datasets.utils import disable_progress_bars, enable_progress_bars
-from diffusers.optimization import get_scheduler
 
 from lerobot.common.datasets.factory import make_dataset
 from lerobot.common.datasets.utils import cycle
@@ -55,6 +54,8 @@ def make_optimizer_and_scheduler(cfg, policy):
             cfg.training.adam_weight_decay,
         )
         assert cfg.training.online_steps == 0, "Diffusion Policy does not handle online training."
+        from diffusers.optimization import get_scheduler
+
         lr_scheduler = get_scheduler(
             cfg.training.lr_scheduler,
             optimizer=optimizer,
@@ -336,7 +337,7 @@ def train(cfg: dict, out_dir=None, job_name=None):
     logging.info(f"{num_total_params=} ({format_big_number(num_total_params)})")
 
     # Note: this helper will be used in offline and online training loops.
-    def _maybe_eval_and_maybe_save(step):
+    def evaluate_and_checkpoint_if_needed(step):
         if step % cfg.training.eval_freq == 0:
             logging.info(f"Eval policy at step {step}")
             eval_info = eval_policy(
@@ -392,9 +393,9 @@ def train(cfg: dict, out_dir=None, job_name=None):
         if step % cfg.training.log_freq == 0:
             log_train_info(logger, train_info, step, cfg, offline_dataset, is_offline)
 
-        # Note: _maybe_eval_and_maybe_save happens **after** the `step`th training update has completed, so we pass in
-        # step + 1.
-        _maybe_eval_and_maybe_save(step + 1)
+        # Note: evaluate_and_checkpoint_if_needed happens **after** the `step`th training update has completed,
+        # so we pass in step + 1.
+        evaluate_and_checkpoint_if_needed(step + 1)
 
         step += 1
 
@@ -460,9 +461,9 @@ def train(cfg: dict, out_dir=None, job_name=None):
             if step % cfg.training.log_freq == 0:
                 log_train_info(logger, train_info, step, cfg, online_dataset, is_offline)
 
-            # Note: _maybe_eval_and_maybe_save happens **after** the `step`th training update has completed, so we pass
-            # in step + 1.
-            _maybe_eval_and_maybe_save(step + 1)
+            # Note: evaluate_and_checkpoint_if_needed happens **after** the `step`th training update has completed,
+            # so we pass in step + 1.
+            evaluate_and_checkpoint_if_needed(step + 1)
 
             step += 1
             online_step += 1

lerobot/scripts/visualize_dataset.py

@@ -32,7 +32,7 @@ local$ rerun lerobot_pusht_episode_0.rrd
 ```
 
 - Visualize data stored on a distant machine through streaming:
-(You need to forward the websocket port to the distant machine, with 
+(You need to forward the websocket port to the distant machine, with
 `ssh -L 9087:localhost:9087 username@remote-host`)
 ```
 distant$ python lerobot/scripts/visualize_dataset.py \
@@ -131,7 +131,7 @@ def visualize_dataset(
             rr.set_time_seconds("timestamp", batch["timestamp"][i].item())
 
             # display each camera image
-            for key in dataset.image_keys:
+            for key in dataset.camera_keys:
                 # TODO(rcadene): add `.compress()`? is it lossless?
                 rr.log(key, rr.Image(to_hwc_uint8_numpy(batch[key][i])))