mirror of
https://github.com/huggingface/lerobot.git
synced 2026-05-15 16:49:55 +00:00
Steven Palma
e881fb6678
* refactor(processor): signature of transform_features * refactor(processor): remove prefixes + processor respect new transform_features signature + update test accordingly * refactor(processor): rename now is only for visual * refactor(processor): update normalize processor * refactor(processor): update vanilla processor features * refactor(processor): feature contract now uses its own enum * chore(processor): rename renameprocessor * chore(processor): minor changes * refactor(processor): add create & change aggregate * refactor(processor): update aggregate * refactor(processor): simplify to functions, fix features contracts and rename function * test(processor): remove to converter tests as now they are very simple * chore(docs): recover docs joint observations processor * fix(processor): update RKP * fix(tests): recv diff test_pipeline * chore(tests): add docs to test * chore(processor): leave obs language constant untouched * fix(processor): correct new shape of feature in crop image processor
398 lines
14 KiB
Python
398 lines
14 KiB
Python
import numpy as np
|
|
import pytest
|
|
import torch
|
|
|
|
from lerobot.processor import TransitionKey
|
|
from lerobot.processor.converters import (
|
|
batch_to_transition,
|
|
to_tensor,
|
|
transition_to_batch,
|
|
transition_to_dataset_frame,
|
|
)
|
|
|
|
|
|
def test_transition_to_dataset_frame_merge_and_pack_vectors_and_metadata():
|
|
# Fabricate dataset features (as stored in dataset.meta["features"])
|
|
features = {
|
|
# Action vector: 3 elements in specific order
|
|
"action": {
|
|
"dtype": "float32",
|
|
"shape": (3,),
|
|
"names": ["j1.pos", "j2.pos", "gripper.pos"],
|
|
},
|
|
# Observation state vector: 2 elements
|
|
"observation.state": {
|
|
"dtype": "float32",
|
|
"shape": (2,),
|
|
"names": ["j1.pos", "j2.pos"],
|
|
},
|
|
# Image spec (video/image dtype acceptable)
|
|
"observation.images.front": {
|
|
"dtype": "image",
|
|
"shape": (480, 640, 3),
|
|
"names": ["h", "w", "c"],
|
|
},
|
|
}
|
|
|
|
# Build two transitions to be merged: teleop (action) and robot obs (state/images)
|
|
img = np.random.randint(0, 255, size=(480, 640, 3), dtype=np.uint8)
|
|
|
|
teleop_transition = {
|
|
TransitionKey.OBSERVATION: {},
|
|
TransitionKey.ACTION: {
|
|
"action.j1.pos": torch.tensor(1.1),
|
|
"action.j2.pos": torch.tensor(2.2),
|
|
# gripper.pos missing → defaults to 0.0
|
|
"action.ee.x": 0.5, # ignored, not in features["action"]["names"]
|
|
},
|
|
TransitionKey.COMPLEMENTARY_DATA: {
|
|
"frame_is_pad": True,
|
|
"task": "Pick cube",
|
|
},
|
|
}
|
|
|
|
robot_transition = {
|
|
TransitionKey.OBSERVATION: {
|
|
"observation.state.j1.pos": torch.tensor(10.0),
|
|
"observation.state.j2.pos": torch.tensor(20.0),
|
|
"observation.images.front": img,
|
|
},
|
|
TransitionKey.REWARD: torch.tensor(5.0),
|
|
TransitionKey.DONE: True,
|
|
TransitionKey.TRUNCATED: False,
|
|
TransitionKey.INFO: {"note": "ok"},
|
|
}
|
|
|
|
# Directly call the refactored function
|
|
batch = transition_to_dataset_frame([teleop_transition, robot_transition], features)
|
|
|
|
# Images passthrough
|
|
assert "observation.images.front" in batch
|
|
assert batch["observation.images.front"].shape == img.shape
|
|
assert batch["observation.images.front"].dtype == np.uint8
|
|
assert np.shares_memory(batch["observation.images.front"], img) or np.array_equal(
|
|
batch["observation.images.front"], img
|
|
)
|
|
|
|
# Observation.state vector
|
|
assert "observation.state" in batch
|
|
obs_vec = batch["observation.state"]
|
|
assert isinstance(obs_vec, np.ndarray) and obs_vec.dtype == np.float32
|
|
assert obs_vec.shape == (2,)
|
|
assert obs_vec[0] == pytest.approx(10.0)
|
|
assert obs_vec[1] == pytest.approx(20.0)
|
|
|
|
# Action vector
|
|
assert "action" in batch
|
|
act_vec = batch["action"]
|
|
assert isinstance(act_vec, np.ndarray) and act_vec.dtype == np.float32
|
|
assert act_vec.shape == (3,)
|
|
assert act_vec[0] == pytest.approx(1.1)
|
|
assert act_vec[1] == pytest.approx(2.2)
|
|
assert act_vec[2] == pytest.approx(0.0) # default for missing gripper.pos
|
|
|
|
# Next.* metadata
|
|
assert batch["next.reward"] == pytest.approx(5.0)
|
|
assert batch["next.done"] is True
|
|
assert batch["next.truncated"] is False
|
|
|
|
# Complementary data
|
|
assert batch["frame_is_pad"] is True
|
|
assert batch["task"] == "Pick cube"
|
|
|
|
|
|
# Tests for the unified to_tensor function
|
|
def test_to_tensor_numpy_arrays():
|
|
"""Test to_tensor with various numpy arrays."""
|
|
# Regular numpy array
|
|
arr = np.array([1.0, 2.0, 3.0])
|
|
result = to_tensor(arr)
|
|
assert isinstance(result, torch.Tensor)
|
|
assert result.dtype == torch.float32
|
|
assert torch.allclose(result, torch.tensor([1.0, 2.0, 3.0]))
|
|
|
|
# Different numpy dtypes should convert to float32 by default
|
|
int_arr = np.array([1, 2, 3], dtype=np.int64)
|
|
result = to_tensor(int_arr)
|
|
assert isinstance(result, torch.Tensor)
|
|
assert result.dtype == torch.float32
|
|
assert torch.allclose(result, torch.tensor([1.0, 2.0, 3.0]))
|
|
|
|
# uint8 arrays (previously "preserved") should now convert
|
|
uint8_arr = np.array([100, 150, 200], dtype=np.uint8)
|
|
result = to_tensor(uint8_arr)
|
|
assert isinstance(result, torch.Tensor)
|
|
assert result.dtype == torch.float32
|
|
assert torch.allclose(result, torch.tensor([100.0, 150.0, 200.0]))
|
|
|
|
|
|
def test_to_tensor_numpy_scalars():
|
|
"""Test to_tensor with numpy scalars (0-dimensional arrays)."""
|
|
# numpy float32 scalar
|
|
scalar = np.float32(3.14)
|
|
result = to_tensor(scalar)
|
|
assert isinstance(result, torch.Tensor)
|
|
assert result.ndim == 0 # Should be 0-dimensional tensor
|
|
assert result.dtype == torch.float32
|
|
assert result.item() == pytest.approx(3.14)
|
|
|
|
# numpy int32 scalar
|
|
int_scalar = np.int32(42)
|
|
result = to_tensor(int_scalar)
|
|
assert isinstance(result, torch.Tensor)
|
|
assert result.ndim == 0
|
|
assert result.dtype == torch.float32
|
|
assert result.item() == pytest.approx(42.0)
|
|
|
|
|
|
def test_to_tensor_python_scalars():
|
|
"""Test to_tensor with Python scalars."""
|
|
# Python int
|
|
result = to_tensor(42)
|
|
assert isinstance(result, torch.Tensor)
|
|
assert result.dtype == torch.float32
|
|
assert result.item() == pytest.approx(42.0)
|
|
|
|
# Python float
|
|
result = to_tensor(3.14)
|
|
assert isinstance(result, torch.Tensor)
|
|
assert result.dtype == torch.float32
|
|
assert result.item() == pytest.approx(3.14)
|
|
|
|
|
|
def test_to_tensor_sequences():
|
|
"""Test to_tensor with lists and tuples."""
|
|
# List
|
|
result = to_tensor([1, 2, 3])
|
|
assert isinstance(result, torch.Tensor)
|
|
assert result.dtype == torch.float32
|
|
assert torch.allclose(result, torch.tensor([1.0, 2.0, 3.0]))
|
|
|
|
# Tuple
|
|
result = to_tensor((4.5, 5.5, 6.5))
|
|
assert isinstance(result, torch.Tensor)
|
|
assert result.dtype == torch.float32
|
|
assert torch.allclose(result, torch.tensor([4.5, 5.5, 6.5]))
|
|
|
|
|
|
def test_to_tensor_existing_tensors():
|
|
"""Test to_tensor with existing PyTorch tensors."""
|
|
# Tensor with same dtype should pass through with potential device change
|
|
tensor = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32)
|
|
result = to_tensor(tensor)
|
|
assert isinstance(result, torch.Tensor)
|
|
assert result.dtype == torch.float32
|
|
assert torch.allclose(result, tensor)
|
|
|
|
# Tensor with different dtype should convert
|
|
int_tensor = torch.tensor([1, 2, 3], dtype=torch.int64)
|
|
result = to_tensor(int_tensor)
|
|
assert isinstance(result, torch.Tensor)
|
|
assert result.dtype == torch.float32
|
|
assert torch.allclose(result, torch.tensor([1.0, 2.0, 3.0]))
|
|
|
|
|
|
def test_to_tensor_dictionaries():
|
|
"""Test to_tensor with nested dictionaries."""
|
|
# Simple dictionary
|
|
data = {"mean": [0.1, 0.2], "std": np.array([1.0, 2.0]), "count": 42}
|
|
result = to_tensor(data)
|
|
assert isinstance(result, dict)
|
|
assert isinstance(result["mean"], torch.Tensor)
|
|
assert isinstance(result["std"], torch.Tensor)
|
|
assert isinstance(result["count"], torch.Tensor)
|
|
assert torch.allclose(result["mean"], torch.tensor([0.1, 0.2]))
|
|
assert torch.allclose(result["std"], torch.tensor([1.0, 2.0]))
|
|
assert result["count"].item() == pytest.approx(42.0)
|
|
|
|
# Nested dictionary
|
|
nested = {
|
|
"action": {"mean": [0.1, 0.2], "std": [1.0, 2.0]},
|
|
"observation": {"mean": np.array([0.5, 0.6]), "count": 10},
|
|
}
|
|
result = to_tensor(nested)
|
|
assert isinstance(result, dict)
|
|
assert isinstance(result["action"], dict)
|
|
assert isinstance(result["observation"], dict)
|
|
assert isinstance(result["action"]["mean"], torch.Tensor)
|
|
assert isinstance(result["observation"]["mean"], torch.Tensor)
|
|
assert torch.allclose(result["action"]["mean"], torch.tensor([0.1, 0.2]))
|
|
assert torch.allclose(result["observation"]["mean"], torch.tensor([0.5, 0.6]))
|
|
|
|
|
|
def test_to_tensor_none_filtering():
|
|
"""Test that None values are filtered out from dictionaries."""
|
|
data = {"valid": [1, 2, 3], "none_value": None, "nested": {"valid": [4, 5], "also_none": None}}
|
|
result = to_tensor(data)
|
|
assert "none_value" not in result
|
|
assert "also_none" not in result["nested"]
|
|
assert "valid" in result
|
|
assert "valid" in result["nested"]
|
|
assert torch.allclose(result["valid"], torch.tensor([1.0, 2.0, 3.0]))
|
|
|
|
|
|
def test_to_tensor_dtype_parameter():
|
|
"""Test to_tensor with different dtype parameters."""
|
|
arr = np.array([1, 2, 3])
|
|
|
|
# Default dtype (float32)
|
|
result = to_tensor(arr)
|
|
assert result.dtype == torch.float32
|
|
|
|
# Explicit float32
|
|
result = to_tensor(arr, dtype=torch.float32)
|
|
assert result.dtype == torch.float32
|
|
|
|
# Float64
|
|
result = to_tensor(arr, dtype=torch.float64)
|
|
assert result.dtype == torch.float64
|
|
|
|
# Preserve original dtype
|
|
float64_arr = np.array([1.0, 2.0, 3.0], dtype=np.float64)
|
|
result = to_tensor(float64_arr, dtype=None)
|
|
assert result.dtype == torch.float64
|
|
|
|
|
|
def test_to_tensor_device_parameter():
|
|
"""Test to_tensor with device parameter."""
|
|
arr = np.array([1.0, 2.0, 3.0])
|
|
|
|
# CPU device (default)
|
|
result = to_tensor(arr, device="cpu")
|
|
assert result.device.type == "cpu"
|
|
|
|
# CUDA device (if available)
|
|
if torch.cuda.is_available():
|
|
result = to_tensor(arr, device="cuda")
|
|
assert result.device.type == "cuda"
|
|
|
|
|
|
def test_to_tensor_empty_dict():
|
|
"""Test to_tensor with empty dictionary."""
|
|
result = to_tensor({})
|
|
assert isinstance(result, dict)
|
|
assert len(result) == 0
|
|
|
|
|
|
def test_to_tensor_unsupported_type():
|
|
"""Test to_tensor with unsupported types raises TypeError."""
|
|
with pytest.raises(TypeError, match="Unsupported type for tensor conversion"):
|
|
to_tensor("unsupported_string")
|
|
|
|
with pytest.raises(TypeError, match="Unsupported type for tensor conversion"):
|
|
to_tensor(object())
|
|
|
|
|
|
def create_transition(
|
|
observation=None, action=None, reward=0.0, done=False, truncated=False, info=None, complementary_data=None
|
|
):
|
|
"""Helper to create an EnvTransition dictionary."""
|
|
return {
|
|
TransitionKey.OBSERVATION: observation,
|
|
TransitionKey.ACTION: action,
|
|
TransitionKey.REWARD: reward,
|
|
TransitionKey.DONE: done,
|
|
TransitionKey.TRUNCATED: truncated,
|
|
TransitionKey.INFO: info if info is not None else {},
|
|
TransitionKey.COMPLEMENTARY_DATA: complementary_data if complementary_data is not None else {},
|
|
}
|
|
|
|
|
|
def test_batch_to_transition_with_index_fields():
|
|
"""Test that batch_to_transition handles index and task_index fields correctly."""
|
|
|
|
# Create batch with index and task_index fields
|
|
batch = {
|
|
"observation.state": torch.randn(1, 7),
|
|
"action": torch.randn(1, 4),
|
|
"next.reward": 1.5,
|
|
"next.done": False,
|
|
"task": ["pick_cube"],
|
|
"index": torch.tensor([42], dtype=torch.int64),
|
|
"task_index": torch.tensor([3], dtype=torch.int64),
|
|
}
|
|
|
|
transition = batch_to_transition(batch)
|
|
|
|
# Check basic transition structure
|
|
assert TransitionKey.OBSERVATION in transition
|
|
assert TransitionKey.ACTION in transition
|
|
assert TransitionKey.COMPLEMENTARY_DATA in transition
|
|
|
|
# Check that index and task_index are in complementary_data
|
|
comp_data = transition[TransitionKey.COMPLEMENTARY_DATA]
|
|
assert "index" in comp_data
|
|
assert "task_index" in comp_data
|
|
assert "task" in comp_data
|
|
|
|
# Verify values
|
|
assert torch.equal(comp_data["index"], batch["index"])
|
|
assert torch.equal(comp_data["task_index"], batch["task_index"])
|
|
assert comp_data["task"] == batch["task"]
|
|
|
|
|
|
def testtransition_to_batch_with_index_fields():
|
|
"""Test that transition_to_batch handles index and task_index fields correctly."""
|
|
|
|
# Create transition with index and task_index in complementary_data
|
|
transition = create_transition(
|
|
observation={"observation.state": torch.randn(1, 7)},
|
|
action=torch.randn(1, 4),
|
|
reward=1.5,
|
|
done=False,
|
|
complementary_data={
|
|
"task": ["navigate"],
|
|
"index": torch.tensor([100], dtype=torch.int64),
|
|
"task_index": torch.tensor([5], dtype=torch.int64),
|
|
},
|
|
)
|
|
|
|
batch = transition_to_batch(transition)
|
|
|
|
# Check that index and task_index are in the batch
|
|
assert "index" in batch
|
|
assert "task_index" in batch
|
|
assert "task" in batch
|
|
|
|
# Verify values
|
|
assert torch.equal(batch["index"], transition[TransitionKey.COMPLEMENTARY_DATA]["index"])
|
|
assert torch.equal(batch["task_index"], transition[TransitionKey.COMPLEMENTARY_DATA]["task_index"])
|
|
assert batch["task"] == transition[TransitionKey.COMPLEMENTARY_DATA]["task"]
|
|
|
|
|
|
def test_batch_to_transition_without_index_fields():
|
|
"""Test that conversion works without index and task_index fields."""
|
|
|
|
# Batch without index/task_index
|
|
batch = {
|
|
"observation.state": torch.randn(1, 7),
|
|
"action": torch.randn(1, 4),
|
|
"task": ["pick_cube"],
|
|
}
|
|
|
|
transition = batch_to_transition(batch)
|
|
comp_data = transition[TransitionKey.COMPLEMENTARY_DATA]
|
|
|
|
# Should have task but not index/task_index
|
|
assert "task" in comp_data
|
|
assert "index" not in comp_data
|
|
assert "task_index" not in comp_data
|
|
|
|
|
|
def test_transition_to_batch_without_index_fields():
|
|
"""Test that conversion works without index and task_index fields."""
|
|
|
|
# Transition without index/task_index
|
|
transition = create_transition(
|
|
observation={"observation.state": torch.randn(1, 7)},
|
|
action=torch.randn(1, 4),
|
|
complementary_data={"task": ["navigate"]},
|
|
)
|
|
|
|
batch = transition_to_batch(transition)
|
|
|
|
# Should have task but not index/task_index
|
|
assert "task" in batch
|
|
assert "index" not in batch
|
|
assert "task_index" not in batch
|