robocandywrapper 0.2.16__tar.gz → 0.2.18__tar.gz

{robocandywrapper-0.2.16/robocandywrapper.egg-info → robocandywrapper-0.2.18}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: robocandywrapper
-Version: 0.2.16
+Version: 0.2.18
 Summary: Sweet wrappers for extending and remixing LeRobot Datasets
 Author: RoboCandyWrapper Contributors
 License: MIT License
@@ -42,10 +42,6 @@ Requires-Dist: numpy>=1.20.0
 Requires-Dist: torch>=2.0.0
 Requires-Dist: lerobot<0.5,>=0.4
 Requires-Dist: pandas>=1.3.0
-Requires-Dist: rewact_tools
-Requires-Dist: motion-primitives
-Requires-Dist: lerobot-policy-diffusion-motion-primitives
-Requires-Dist: lerobot-policy-diffusion-pointing
 Provides-Extra: dev
 Requires-Dist: pytest>=7.0.0; extra == "dev"
 Requires-Dist: black>=22.0.0; extra == "dev"

{robocandywrapper-0.2.16 → robocandywrapper-0.2.18}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "robocandywrapper"
-version = "0.2.16"
+version = "0.2.18"
 description = "Sweet wrappers for extending and remixing LeRobot Datasets"
 readme = "README.md"
 requires-python = ">=3.10,<3.11"
@@ -27,10 +27,6 @@ dependencies = [
     "torch>=2.0.0",
     "lerobot>=0.4,<0.5",
     "pandas>=1.3.0",
-    "rewact_tools",
-    "motion-primitives",
-    "lerobot-policy-diffusion-motion-primitives",
-    "lerobot-policy-diffusion-pointing",
 ]
 
 [project.optional-dependencies]
@@ -56,12 +52,6 @@ override-dependencies = [
     "evdev>=0.0.0; sys_platform == 'never'",
 ]
 
-[tool.uv.sources]
-lerobot-policy-diffusion-motion-primitives = { path = "../../experimental/lerobot_policy_diffusion_motion_primitives", editable = true }
-lerobot-policy-diffusion-pointing = { path = "../../experimental/lerobot_policy_diffusion_pointing", editable = true }
-motion-primitives = { path = "../../experimental/motion_primitives", editable = true }
-rewact-tools = { path = "../rewACT/rewact_tools", editable = true }
-
 [tool.black]
 line-length = 100
 target-version = ['py310']

{robocandywrapper-0.2.16 → robocandywrapper-0.2.18}/robocandywrapper/__init__.py

@@ -18,7 +18,7 @@ from robocandywrapper.constants import (
     MOTION_PRIMITIVE_PLUGIN_NAME,
 )
 
-__version__ = "0.2.16"
+__version__ = "0.2.18"
 
 __all__ = [
     "DatasetPlugin",

{robocandywrapper-0.2.16 → robocandywrapper-0.2.18}/robocandywrapper/metadata_view.py

@@ -8,21 +8,22 @@ import numpy as np
 import pandas as pd
 
 
-def _task_names_from_meta_tasks(meta_tasks) -> set[str]:
+def _task_index_to_name(meta_tasks) -> dict[int, str]:
     """
-    Extract task name strings from a dataset's meta.tasks in either format.
+    Normalize a dataset's task metadata to task_index -> task name.
 
     - LeRobot 2.1 (dict): tasks is dict[task_index, task_name].
     - LeRobot 3.0 (parquet/DataFrame): tasks is a pandas DataFrame where the
       *row index* is the task name (string) and the only column is "task_index"
       (integer). So: df.index = task names, df["task_index"] = indices.
 
-    Returns a set of unique task name strings.
+    Some task tables instead include both "task_index" and "task" columns.
     """
     if meta_tasks is None:
-        return set()
+        return {}
     if isinstance(meta_tasks, dict):
-        return {str(v) for v in meta_tasks.values() if isinstance(v, str)}
+        return {int(task_index): str(name) for task_index, name in meta_tasks.items()}
+
     # DataFrame or object with to_pandas() (e.g. HuggingFace Dataset)
     df = None
     if hasattr(meta_tasks, "to_pandas"):
@@ -30,14 +31,24 @@ def _task_names_from_meta_tasks(meta_tasks) -> set[str]:
     elif isinstance(meta_tasks, pd.DataFrame):
         df = meta_tasks
     if df is None:
-        return set()
-    # 3.0 format: optional "task" column, else task names are the row index
-    if "task" in df.columns:
-        return set(df["task"].dropna().astype(str).unique())
+        return {}
+
     if "task_index" in df.columns:
-        # Task names are the DataFrame index (row labels), not the column
-        return set(df.index.astype(str))
-    return set()
+        if "task" in df.columns:
+            return {
+                int(row["task_index"]): str(row["task"])
+                for _, row in df.iterrows()
+            }
+        return {
+            int(row["task_index"]): str(task_name)
+            for task_name, row in df.iterrows()
+        }
+    if "task" in df.columns:
+        return {
+            int(task_index): str(name)
+            for task_index, name in enumerate(df["task"])
+        }
+    return {}
 
 
 def aggregate_stats_weighted(
@@ -301,7 +312,7 @@ class WrappedRobotDatasetMetadataView:
         all_task_names: set[str] = set()
         for dataset in self._datasets:
             meta_tasks = getattr(dataset.meta, "tasks", None)
-            all_task_names.update(_task_names_from_meta_tasks(meta_tasks))
+            all_task_names.update(_task_index_to_name(meta_tasks).values())
 
         sorted_names = sorted(all_task_names)
         self._tasks = {idx: name for idx, name in enumerate(sorted_names)}

{robocandywrapper-0.2.16 → robocandywrapper-0.2.18}/robocandywrapper/wrapper.py

@@ -9,7 +9,10 @@ from lerobot.configs.types import FeatureType, PolicyFeature
 import torch
 
 from robocandywrapper import DatasetPlugin, PluginConflictError, PluginInstance
-from robocandywrapper.metadata_view import WrappedRobotDatasetMetadataView
+from robocandywrapper.metadata_view import (
+    WrappedRobotDatasetMetadataView,
+    _task_index_to_name,
+)
 
 
 class WrappedRobotDataset(torch.utils.data.Dataset):
@@ -114,6 +117,15 @@ class WrappedRobotDataset(torch.utils.data.Dataset):
             dataset_weights=dataset_weights,
             dataset_renames=self._dataset_renames,
         )
+        self._task_index_remaps = [
+            {
+                inner_index: self._meta.task_to_task_index[task_name]
+                for inner_index, task_name in _task_index_to_name(
+                    getattr(dataset.meta, "tasks", None)
+                ).items()
+            }
+            for dataset in self._datasets
+        ]
 
         # ** MATCHING LeRobot MULTI-DATASET API DESIGN **
         
@@ -627,6 +639,25 @@ class WrappedRobotDataset(torch.utils.data.Dataset):
         else:
             item = dataset[local_idx]
 
+        if "task_index" in item and self._task_index_remaps[dataset_idx]:
+            task_index = item["task_index"]
+            inner_task_index = (
+                int(task_index.item()) if isinstance(task_index, torch.Tensor) else int(task_index)
+            )
+            try:
+                unified_task_index = self._task_index_remaps[dataset_idx][inner_task_index]
+            except KeyError:
+                raise KeyError(
+                    f"task_index {inner_task_index} from dataset {dataset.repo_id!r} "
+                    "is missing from its task metadata"
+                ) from None
+
+            if isinstance(task_index, torch.Tensor):
+                item["task_index"] = task_index.new_tensor(unified_task_index)
+            else:
+                item["task_index"] = unified_task_index
+            item["task"] = self._meta.tasks[unified_task_index]
+
         episode_idx = item["episode_index"].item() if isinstance(item["episode_index"], torch.Tensor) else item["episode_index"]
         
         # Add dataset index

{robocandywrapper-0.2.16 → robocandywrapper-0.2.18/robocandywrapper.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: robocandywrapper
-Version: 0.2.16
+Version: 0.2.18
 Summary: Sweet wrappers for extending and remixing LeRobot Datasets
 Author: RoboCandyWrapper Contributors
 License: MIT License
@@ -42,10 +42,6 @@ Requires-Dist: numpy>=1.20.0
 Requires-Dist: torch>=2.0.0
 Requires-Dist: lerobot<0.5,>=0.4
 Requires-Dist: pandas>=1.3.0
-Requires-Dist: rewact_tools
-Requires-Dist: motion-primitives
-Requires-Dist: lerobot-policy-diffusion-motion-primitives
-Requires-Dist: lerobot-policy-diffusion-pointing
 Provides-Extra: dev
 Requires-Dist: pytest>=7.0.0; extra == "dev"
 Requires-Dist: black>=22.0.0; extra == "dev"

{robocandywrapper-0.2.16 → robocandywrapper-0.2.18}/robocandywrapper.egg-info/SOURCES.txt

@@ -31,4 +31,5 @@ robocandywrapper/samplers/config.py
 robocandywrapper/samplers/factory.py
 robocandywrapper/samplers/weighted.py
 tests/test_dataset_weights_integration.py
-tests/test_key_rename_stats.py
+tests/test_key_rename_stats.py
+tests/test_task_index_remapping.py

{robocandywrapper-0.2.16 → robocandywrapper-0.2.18}/robocandywrapper.egg-info/requires.txt

@@ -2,10 +2,6 @@ numpy>=1.20.0
 torch>=2.0.0
 lerobot<0.5,>=0.4
 pandas>=1.3.0
-rewact_tools
-motion-primitives
-lerobot-policy-diffusion-motion-primitives
-lerobot-policy-diffusion-pointing
 
 [dev]
 pytest>=7.0.0

{robocandywrapper-0.2.16 → robocandywrapper-0.2.18}/setup.py

@@ -9,7 +9,7 @@ long_description = readme_file.read_text(encoding="utf-8") if readme_file.exists
 
 setup(
     name="robocandywrapper",
-    version="0.2.16",
+    version="0.2.18",
     description="Sweet wrappers for extending and remixing LeRobot Datasets",
     long_description=long_description,
     long_description_content_type="text/markdown",

robocandywrapper-0.2.18/tests/test_task_index_remapping.py

@@ -0,0 +1,134 @@
+import sys
+
+import lerobot.datasets.utils as lerobot_dataset_utils
+import pandas as pd
+import pytest
+import torch
+
+# Current main imports these helpers from the newer module location. Alias the
+# older location when running against the repository's declared LeRobot range.
+sys.modules.setdefault("lerobot.datasets.feature_utils", lerobot_dataset_utils)
+sys.modules.setdefault("lerobot.datasets.io_utils", lerobot_dataset_utils)
+
+from robocandywrapper.wrapper import WrappedRobotDataset
+
+
+class PandasTasks:
+    def __init__(self, index_to_task: dict[int, str]):
+        self._frame = pd.DataFrame(
+            {"task_index": list(index_to_task)},
+            index=pd.Index(index_to_task.values(), name="task"),
+        )
+
+    def to_pandas(self) -> pd.DataFrame:
+        return self._frame.copy()
+
+
+class MockMetadata:
+    def __init__(self, tasks):
+        self.tasks = tasks
+        self.stats = {}
+        self.info = {"fps": 20, "video": False}
+        self.camera_keys = []
+        self.image_keys = []
+        self.video_keys = []
+        self.total_frames = 1
+        self.total_episodes = 1
+        self.episodes = {}
+
+    @property
+    def fps(self):
+        return self.info["fps"]
+
+    @property
+    def features(self):
+        return {
+            "episode_index": {"dtype": "int64", "shape": [1]},
+            "task_index": {"dtype": "int64", "shape": [1]},
+        }
+
+
+class MockDataset:
+    def __init__(self, repo_id: str, tasks, task_index):
+        self.repo_id = repo_id
+        self.meta = MockMetadata(tasks)
+        self.features = self.meta.features
+        self.hf_features = self.features
+        self._task_index = task_index
+        self.episodes = None
+        self.num_episodes = 1
+
+    def __len__(self):
+        return 1
+
+    def __getitem__(self, index):
+        return {
+            "episode_index": torch.tensor(0),
+            "task_index": (
+                self._task_index.clone()
+                if isinstance(self._task_index, torch.Tensor)
+                else self._task_index
+            ),
+        }
+
+
+def test_single_dataset_sample_uses_sorted_wrapper_task_index():
+    dataset = MockDataset(
+        "single",
+        PandasTasks({0: "z task", 1: "a task"}),
+        torch.tensor(0, dtype=torch.int64),
+    )
+
+    wrapped = WrappedRobotDataset(dataset)
+    item = wrapped[0]
+
+    assert wrapped.meta.tasks == {0: "a task", 1: "z task"}
+    assert torch.equal(item["task_index"], torch.tensor(1, dtype=torch.int64))
+    assert item["task"] == "z task"
+
+
+def test_mixed_dataset_samples_remap_colliding_local_task_indices():
+    first = MockDataset("first", {0: "z task"}, 0)
+    second = MockDataset("second", PandasTasks({0: "a task"}), 0)
+
+    wrapped = WrappedRobotDataset([first, second])
+
+    assert wrapped.meta.tasks == {0: "a task", 1: "z task"}
+    assert wrapped[0]["task_index"] == 1
+    assert wrapped[0]["task"] == "z task"
+    assert wrapped[1]["task_index"] == 0
+    assert wrapped[1]["task"] == "a task"
+
+
+def test_unknown_inner_task_index_has_contextual_error():
+    dataset = MockDataset("broken", {0: "known task"}, 9)
+    wrapped = WrappedRobotDataset(dataset)
+
+    with pytest.raises(KeyError, match=r"task_index 9.*broken"):
+        wrapped[0]
+
+
+def test_real_busybox_sample_remaps_task_index_without_changing_task(monkeypatch):
+    from lerobot.datasets.lerobot_dataset import LeRobotDataset
+
+    inner = LeRobotDataset(
+        "villekuosmanen/busybox_multitask",
+        episodes=[12],
+        download_videos=False,
+    )
+    monkeypatch.setattr(inner, "_query_videos", lambda *args, **kwargs: {})
+
+    inner_item = inner[0]
+    inner_task_index = int(inner_item["task_index"])
+    inner_task = str(inner.meta.tasks.index[inner.meta.tasks["task_index"] == inner_task_index][0])
+
+    wrapped = WrappedRobotDataset(inner)
+    wrapped_item = wrapped[0]
+    wrapped_task_index = int(wrapped_item["task_index"])
+
+    assert inner_task_index == 4
+    assert inner_task == "Move the right slider to position 5"
+    assert wrapped.meta.tasks[4] == "Move the left slider to position 5"
+    assert wrapped_task_index == 9
+    assert wrapped.meta.tasks[wrapped_task_index] == inner_task
+    assert wrapped_item["task"] == inner_task