chore: import upstream snapshot with attribution

2026-07-13 12:58:47 +08:00
commit e68525745e
990 changed files with 259649 additions and 0 deletions
@@ -0,0 +1,296 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+
+from lerobot.processor import DataProcessorPipeline
+from lerobot.processor.converters import batch_to_transition, transition_to_batch
+from lerobot.types import TransitionKey
+from lerobot.utils.constants import ACTION, DONE, OBS_IMAGE, OBS_PREFIX, OBS_STATE, REWARD, TRUNCATED
+
+
+def _dummy_batch():
+    """Create a dummy batch using the new format with observation.* and next.* keys."""
+    return {
+        f"{OBS_IMAGE}.left": torch.randn(1, 3, 128, 128),
+        f"{OBS_IMAGE}.right": torch.randn(1, 3, 128, 128),
+        OBS_STATE: torch.tensor([[0.1, 0.2, 0.3, 0.4]]),
+        ACTION: torch.tensor([[0.5]]),
+        REWARD: 1.0,
+        DONE: False,
+        TRUNCATED: False,
+        "info": {"key": "value"},
+    }
+
+
+def test_observation_grouping_roundtrip():
+    """Test that observation.* keys are properly grouped and ungrouped."""
+    proc = DataProcessorPipeline([])
+    batch_in = _dummy_batch()
+    batch_out = proc(batch_in)
+
+    # Check that all observation.* keys are preserved
+    original_obs_keys = {k: v for k, v in batch_in.items() if k.startswith(OBS_PREFIX)}
+    reconstructed_obs_keys = {k: v for k, v in batch_out.items() if k.startswith(OBS_PREFIX)}
+
+    assert set(original_obs_keys.keys()) == set(reconstructed_obs_keys.keys())
+
+    # Check tensor values
+    assert torch.allclose(batch_out[f"{OBS_IMAGE}.left"], batch_in[f"{OBS_IMAGE}.left"])
+    assert torch.allclose(batch_out[f"{OBS_IMAGE}.right"], batch_in[f"{OBS_IMAGE}.right"])
+    assert torch.allclose(batch_out[OBS_STATE], batch_in[OBS_STATE])
+
+    # Check other fields
+    assert torch.allclose(batch_out[ACTION], batch_in[ACTION])
+    assert batch_out[REWARD] == batch_in[REWARD]
+    assert batch_out[DONE] == batch_in[DONE]
+    assert batch_out[TRUNCATED] == batch_in[TRUNCATED]
+    assert batch_out["info"] == batch_in["info"]
+
+
+def test_batch_to_transition_observation_grouping():
+    """Test that batch_to_transition correctly groups observation.* keys."""
+    batch = {
+        f"{OBS_IMAGE}.top": torch.randn(1, 3, 128, 128),
+        f"{OBS_IMAGE}.left": torch.randn(1, 3, 128, 128),
+        OBS_STATE: [1, 2, 3, 4],
+        ACTION: torch.tensor([0.1, 0.2, 0.3, 0.4]),
+        REWARD: 1.5,
+        DONE: True,
+        TRUNCATED: False,
+        "info": {"episode": 42},
+    }
+
+    transition = batch_to_transition(batch)
+
+    # Check observation is a dict with all observation.* keys
+    assert isinstance(transition[TransitionKey.OBSERVATION], dict)
+    assert f"{OBS_IMAGE}.top" in transition[TransitionKey.OBSERVATION]
+    assert f"{OBS_IMAGE}.left" in transition[TransitionKey.OBSERVATION]
+    assert OBS_STATE in transition[TransitionKey.OBSERVATION]
+
+    # Check values are preserved
+    assert torch.allclose(
+        transition[TransitionKey.OBSERVATION][f"{OBS_IMAGE}.top"], batch[f"{OBS_IMAGE}.top"]
+    )
+    assert torch.allclose(
+        transition[TransitionKey.OBSERVATION][f"{OBS_IMAGE}.left"], batch[f"{OBS_IMAGE}.left"]
+    )
+    assert transition[TransitionKey.OBSERVATION][OBS_STATE] == [1, 2, 3, 4]
+
+    # Check other fields
+    assert torch.allclose(transition[TransitionKey.ACTION], torch.tensor([0.1, 0.2, 0.3, 0.4]))
+    assert transition[TransitionKey.REWARD] == 1.5
+    assert transition[TransitionKey.DONE]
+    assert not transition[TransitionKey.TRUNCATED]
+    assert transition[TransitionKey.INFO] == {"episode": 42}
+    assert transition[TransitionKey.COMPLEMENTARY_DATA] == {}
+
+
+def test_transition_to_batch_observation_flattening():
+    """Test that transition_to_batch correctly flattens observation dict."""
+    observation_dict = {
+        f"{OBS_IMAGE}.top": torch.randn(1, 3, 128, 128),
+        f"{OBS_IMAGE}.left": torch.randn(1, 3, 128, 128),
+        OBS_STATE: [1, 2, 3, 4],
+    }
+
+    transition = {
+        TransitionKey.OBSERVATION: observation_dict,
+        TransitionKey.ACTION: "action_data",
+        TransitionKey.REWARD: 1.5,
+        TransitionKey.DONE: True,
+        TransitionKey.TRUNCATED: False,
+        TransitionKey.INFO: {"episode": 42},
+        TransitionKey.COMPLEMENTARY_DATA: {},
+    }
+
+    batch = transition_to_batch(transition)
+
+    # Check that observation.* keys are flattened back to batch
+    assert f"{OBS_IMAGE}.top" in batch
+    assert f"{OBS_IMAGE}.left" in batch
+    assert OBS_STATE in batch
+
+    # Check values are preserved
+    assert torch.allclose(batch[f"{OBS_IMAGE}.top"], observation_dict[f"{OBS_IMAGE}.top"])
+    assert torch.allclose(batch[f"{OBS_IMAGE}.left"], observation_dict[f"{OBS_IMAGE}.left"])
+    assert batch[OBS_STATE] == [1, 2, 3, 4]
+
+    # Check other fields are mapped to next.* format
+    assert batch[ACTION] == "action_data"
+    assert batch[REWARD] == 1.5
+    assert batch[DONE]
+    assert not batch[TRUNCATED]
+    assert batch["info"] == {"episode": 42}
+
+
+def test_no_observation_keys():
+    """Test behavior when there are no observation.* keys."""
+    batch = {
+        ACTION: torch.tensor([1.0, 2.0]),
+        REWARD: 2.0,
+        DONE: False,
+        TRUNCATED: True,
+        "info": {"test": "no_obs"},
+    }
+
+    transition = batch_to_transition(batch)
+
+    # Observation should be None when no observation.* keys
+    assert transition[TransitionKey.OBSERVATION] is None
+
+    # Check other fields
+    assert torch.allclose(transition[TransitionKey.ACTION], torch.tensor([1.0, 2.0]))
+    assert transition[TransitionKey.REWARD] == 2.0
+    assert not transition[TransitionKey.DONE]
+    assert transition[TransitionKey.TRUNCATED]
+    assert transition[TransitionKey.INFO] == {"test": "no_obs"}
+
+    # Round trip should work
+    reconstructed_batch = transition_to_batch(transition)
+    assert torch.allclose(reconstructed_batch[ACTION], torch.tensor([1.0, 2.0]))
+    assert reconstructed_batch[REWARD] == 2.0
+    assert not reconstructed_batch[DONE]
+    assert reconstructed_batch[TRUNCATED]
+    assert reconstructed_batch["info"] == {"test": "no_obs"}
+
+
+def test_minimal_batch():
+    """Test with minimal batch containing only observation.* and action."""
+    batch = {OBS_STATE: "minimal_state", ACTION: torch.tensor([0.5])}
+
+    transition = batch_to_transition(batch)
+
+    # Check observation
+    assert transition[TransitionKey.OBSERVATION] == {OBS_STATE: "minimal_state"}
+    assert torch.allclose(transition[TransitionKey.ACTION], torch.tensor([0.5]))
+
+    # Check defaults
+    assert transition[TransitionKey.REWARD] == 0.0
+    assert not transition[TransitionKey.DONE]
+    assert not transition[TransitionKey.TRUNCATED]
+    assert transition[TransitionKey.INFO] == {}
+    assert transition[TransitionKey.COMPLEMENTARY_DATA] == {}
+
+    # Round trip
+    reconstructed_batch = transition_to_batch(transition)
+    assert reconstructed_batch[OBS_STATE] == "minimal_state"
+    assert torch.allclose(reconstructed_batch[ACTION], torch.tensor([0.5]))
+    assert reconstructed_batch[REWARD] == 0.0
+    assert not reconstructed_batch[DONE]
+    assert not reconstructed_batch[TRUNCATED]
+    assert reconstructed_batch["info"] == {}
+
+
+def test_empty_batch():
+    """Test behavior with empty batch."""
+    batch = {}
+
+    transition = batch_to_transition(batch)
+
+    # All fields should have defaults
+    assert transition[TransitionKey.OBSERVATION] is None
+    assert transition[TransitionKey.ACTION] is None
+    assert transition[TransitionKey.REWARD] == 0.0
+    assert not transition[TransitionKey.DONE]
+    assert not transition[TransitionKey.TRUNCATED]
+    assert transition[TransitionKey.INFO] == {}
+    assert transition[TransitionKey.COMPLEMENTARY_DATA] == {}
+
+    # Round trip
+    reconstructed_batch = transition_to_batch(transition)
+    assert reconstructed_batch[ACTION] is None
+    assert reconstructed_batch[REWARD] == 0.0
+    assert not reconstructed_batch[DONE]
+    assert not reconstructed_batch[TRUNCATED]
+    assert reconstructed_batch["info"] == {}
+
+
+def test_complex_nested_observation():
+    """Test with complex nested observation data."""
+    batch = {
+        f"{OBS_IMAGE}.top": {"image": torch.randn(1, 3, 128, 128), "timestamp": 1234567890},
+        f"{OBS_IMAGE}.left": {"image": torch.randn(1, 3, 128, 128), "timestamp": 1234567891},
+        OBS_STATE: torch.randn(7),
+        ACTION: torch.randn(8),
+        REWARD: 3.14,
+        DONE: False,
+        TRUNCATED: True,
+        "info": {"episode_length": 200, "success": True},
+    }
+
+    transition = batch_to_transition(batch)
+    reconstructed_batch = transition_to_batch(transition)
+
+    # Check that all observation keys are preserved
+    original_obs_keys = {k for k in batch if k.startswith(OBS_PREFIX)}
+    reconstructed_obs_keys = {k for k in reconstructed_batch if k.startswith(OBS_PREFIX)}
+
+    assert original_obs_keys == reconstructed_obs_keys
+
+    # Check tensor values
+    assert torch.allclose(batch[OBS_STATE], reconstructed_batch[OBS_STATE])
+
+    # Check nested dict with tensors
+    assert torch.allclose(
+        batch[f"{OBS_IMAGE}.top"]["image"], reconstructed_batch[f"{OBS_IMAGE}.top"]["image"]
+    )
+    assert torch.allclose(
+        batch[f"{OBS_IMAGE}.left"]["image"], reconstructed_batch[f"{OBS_IMAGE}.left"]["image"]
+    )
+
+    # Check action tensor
+    assert torch.allclose(batch[ACTION], reconstructed_batch[ACTION])
+
+    # Check other fields
+    assert batch[REWARD] == reconstructed_batch[REWARD]
+    assert batch[DONE] == reconstructed_batch[DONE]
+    assert batch[TRUNCATED] == reconstructed_batch[TRUNCATED]
+    assert batch["info"] == reconstructed_batch["info"]
+
+
+def test_custom_converter():
+    """Test that custom converters can still be used."""
+
+    def to_tr(batch):
+        # Custom converter that modifies the reward
+        tr = batch_to_transition(batch)
+        # Double the reward
+        reward = tr.get(TransitionKey.REWARD, 0.0)
+        new_tr = tr.copy()
+        new_tr[TransitionKey.REWARD] = reward * 2 if reward is not None else 0.0
+        return new_tr
+
+    def to_batch(tr):
+        batch = transition_to_batch(tr)
+        return batch
+
+    processor = DataProcessorPipeline(steps=[], to_transition=to_tr, to_output=to_batch)
+
+    batch = {
+        OBS_STATE: torch.randn(1, 4),
+        ACTION: torch.randn(1, 2),
+        REWARD: 1.0,
+        DONE: False,
+    }
+
+    result = processor(batch)
+
+    # Check the reward was doubled by our custom converter
+    assert result[REWARD] == 2.0
+    assert torch.allclose(result[OBS_STATE], batch[OBS_STATE])
+    assert torch.allclose(result[ACTION], batch[ACTION])