From 0af46a58db559c84633144bde564f8ad8d33157c Mon Sep 17 00:00:00 2001
From: Remi Cadene <re.cadene@gmail.com>
Date: Fri, 31 May 2024 13:07:41 +0000
Subject: [PATCH] Addressed

---
 lerobot/common/datasets/factory.py            | 22 +++++++++++--------
 .../common/policies/act/configuration_act.py  |  1 +
 2 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/lerobot/common/datasets/factory.py b/lerobot/common/datasets/factory.py
index 7563a806..4732f577 100644
--- a/lerobot/common/datasets/factory.py
+++ b/lerobot/common/datasets/factory.py
@@ -55,15 +55,19 @@ def make_dataset(cfg, split: str = "train") -> LeRobotDataset | MultiLeRobotData
             "strings to load multiple datasets."
         )
 
-    if (
-        cfg.env.name != "dora"
-        and isinstance(cfg.dataset_repo_id, str)
-        and cfg.env.name not in cfg.dataset_repo_id
-    ):
-        logging.warning(
-            f"There might be a mismatch between your training dataset ({cfg.dataset_repo_id=}) and your "
-            f"environment ({cfg.env.name=})."
-        )
+    # A soft check to warn if the environment matches the dataset. Don't check if we are using a real world env (dora).
+    if cfg.env.name != "dora":
+        if isinstance(cfg.dataset_repo_id, str):
+            dataset_repo_ids = [cfg.dataset_repo_id]  # single dataset
+        else:
+            dataset_repo_ids = cfg.dataset_repo_id  # multiple datasets
+
+        for dataset_repo_id in dataset_repo_ids:
+            if cfg.env.name not in dataset_repo_id:
+                logging.warning(
+                    f"There might be a mismatch between your training dataset ({dataset_repo_id=}) and your "
+                    f"environment ({cfg.env.name=})."
+                )
 
     resolve_delta_timestamps(cfg)
 
diff --git a/lerobot/common/policies/act/configuration_act.py b/lerobot/common/policies/act/configuration_act.py
index d33e8bd9..a4b0b7d2 100644
--- a/lerobot/common/policies/act/configuration_act.py
+++ b/lerobot/common/policies/act/configuration_act.py
@@ -29,6 +29,7 @@ class ACTConfig:
         - At least one key starting with "observation.image is required as an input.
         - If there are multiple keys beginning with "observation.images." they are treated as multiple camera
           views. Right now we only support all images having the same shape.
+        - May optionally work without an "observation.state" key for the proprioceptive robot state.
         - "action" is required as an output key.
 
     Args: