From 13441f0d983afd6b2628fd377a0e5fbc05243660 Mon Sep 17 00:00:00 2001
From: Michel Aractingi <michel.aractingi@huggingface.co>
Date: Mon, 30 Dec 2024 13:47:28 +0000
Subject: [PATCH] added temporary fix for missing task_index key in online
 environment

---
 lerobot/scripts/train.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/lerobot/scripts/train.py b/lerobot/scripts/train.py
index fbe7927d..a4eb3528 100644
--- a/lerobot/scripts/train.py
+++ b/lerobot/scripts/train.py
@@ -322,6 +322,11 @@ def train(cfg: DictConfig, out_dir: str | None = None, job_name: str | None = No
 
     logging.info("make_dataset")
     offline_dataset = make_dataset(cfg)
+    # TODO (michel-aractingi): temporary fix to avoid datasets with task_index key that doesn't exist in online environment
+    # i.e., pusht
+    if "task_index" in offline_dataset.hf_dataset[0]:
+        offline_dataset.hf_dataset = offline_dataset.hf_dataset.remove_columns(["task_index"])
+
     if isinstance(offline_dataset, MultiLeRobotDataset):
         logging.info(
             "Multiple datasets were provided. Applied the following index mapping to the provided datasets: "