fix: populate config from dataset metadata

2026-07-25 18:56:09 +00:00 · 2026-07-23 22:17:39 +02:00
parent d015474483
commit 39284f43c7
2 changed files with 55 additions and 0 deletions
@@ -0,0 +1,46 @@
+from types import SimpleNamespace
+
+from torch import nn
+
+from lerobot.configs import FeatureType
+from lerobot.rewards.factory import make_reward_model
+from lerobot.rewards.temporal_siglip_value_function.configuration_temporal_siglip_value_function import (
+    TemporalSiglipVFConfig,
+)
+
+
+def test_reward_factory_populates_input_features_from_dataset_meta(monkeypatch):
+    from lerobot.rewards import factory
+
+    class FakeReward(nn.Module):
+        def __init__(self, config, **kwargs):
+            super().__init__()
+            self.config = config
+
+    monkeypatch.setattr(factory, "get_reward_model_class", lambda name: FakeReward)
+    metadata = SimpleNamespace(
+        features={
+            "observation.images.top": {
+                "dtype": "video",
+                "shape": [480, 640, 3],
+                "names": ["height", "width", "channel"],
+            },
+            "observation.state": {
+                "dtype": "float32",
+                "shape": [14],
+                "names": [f"joint_{index}" for index in range(14)],
+            },
+            "action": {
+                "dtype": "float32",
+                "shape": [14],
+                "names": [f"joint_{index}" for index in range(14)],
+            },
+        }
+    )
+    config = TemporalSiglipVFConfig(device="cpu")
+    model = make_reward_model(config, dataset_meta=metadata)
+
+    assert model.config.input_features["observation.images.top"].type is FeatureType.VISUAL
+    assert model.config.input_features["observation.images.top"].shape == (3, 480, 640)
+    assert model.config.input_features["observation.state"].type is FeatureType.STATE
+    assert "action" not in model.config.input_features