added tests

tests/test_arena_processor.py

@@ -0,0 +1,406 @@
+
+# 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 pytest
+import torch
+
+from lerobot.configs.types import (
+    FeatureType,
+    PipelineFeatureType,
+    PolicyFeature,
+)
+from lerobot.processor.env_processor import IsaaclabArenaProcessorStep
+from lerobot.utils.constants import OBS_IMAGES, OBS_STATE, OBS_STR
+
+# Test constants
+BATCH_SIZE = 2
+STATE_DIM = 16
+IMG_HEIGHT = 64
+IMG_WIDTH = 64
+
+# Generic test keys (not real robot keys)
+TEST_STATE_KEY = "test_state_obs"
+TEST_CAMERA_KEY = "test_rgb_cam"
+
+
+@pytest.fixture
+def processor():
+    """Default processor with test keys."""
+    return IsaaclabArenaProcessorStep(
+        state_keys=(TEST_STATE_KEY,),
+        camera_keys=(TEST_CAMERA_KEY,),
+    )
+
+
+@pytest.fixture
+def sample_observation():
+    """Sample IsaacLab Arena observation with state and camera data."""
+    return {
+        f"{OBS_STR}.policy": {
+            TEST_STATE_KEY: torch.randn(BATCH_SIZE, STATE_DIM),
+        },
+        f"{OBS_STR}.camera_obs": {
+            TEST_CAMERA_KEY: torch.randint(0, 255, (BATCH_SIZE, IMG_HEIGHT, IMG_WIDTH, 3), dtype=torch.uint8),
+        },
+    }
+
+
+# =============================================================================
+# State Processing Tests
+# =============================================================================
+
+
+def test_state_extraction(processor, sample_observation):
+    """Test that state is extracted and converted to float32."""
+    processed = processor.observation(sample_observation)
+
+    assert OBS_STATE in processed
+    assert processed[OBS_STATE].shape == (BATCH_SIZE, STATE_DIM)
+    assert processed[OBS_STATE].dtype == torch.float32
+
+
+def test_state_concatenation_multiple_keys():
+    """Test that multiple state keys are concatenated in order."""
+    dim1, dim2 = 10, 6
+    processor = IsaaclabArenaProcessorStep(
+        state_keys=("state_alpha", "state_beta"),
+        camera_keys=(),
+    )
+
+    obs = {
+        f"{OBS_STR}.policy": {
+            "state_alpha": torch.ones(BATCH_SIZE, dim1),
+            "state_beta": torch.ones(BATCH_SIZE, dim2) * 2,
+        },
+    }
+
+    processed = processor.observation(obs)
+
+    state = processed[OBS_STATE]
+    assert state.shape == (BATCH_SIZE, dim1 + dim2)
+    # Verify ordering: first dim1 elements are 1s, last dim2 are 2s
+    assert torch.all(state[:, :dim1] == 1.0)
+    assert torch.all(state[:, dim1:] == 2.0)
+
+
+def test_state_flattening_higher_dims():
+    """Test that state with dim > 2 is flattened to (B, -1)."""
+    processor = IsaaclabArenaProcessorStep(
+        state_keys=("multidim_state",),
+        camera_keys=(),
+    )
+
+    # Shape (B, 4, 4) -> should flatten to (B, 16)
+    obs = {
+        f"{OBS_STR}.policy": {
+            "multidim_state": torch.randn(BATCH_SIZE, 4, 4),
+        },
+    }
+
+    processed = processor.observation(obs)
+
+    assert processed[OBS_STATE].shape == (BATCH_SIZE, 16)
+
+
+def test_state_filters_to_configured_keys():
+    """Test that only configured state_keys are extracted."""
+    processor = IsaaclabArenaProcessorStep(
+        state_keys=("included_key",),
+        camera_keys=(),
+    )
+
+    obs = {
+        f"{OBS_STR}.policy": {
+            "included_key": torch.randn(BATCH_SIZE, 10),
+            "excluded_key": torch.randn(BATCH_SIZE, 6),  # Should be ignored
+        },
+    }
+
+    processed = processor.observation(obs)
+
+    # Only included_key (dim 10) should be included
+    assert processed[OBS_STATE].shape == (BATCH_SIZE, 10)
+
+
+def test_missing_state_key_skipped():
+    """Test that missing state keys in observation are skipped."""
+    processor = IsaaclabArenaProcessorStep(
+        state_keys=("present_key", "missing_key"),
+        camera_keys=(),
+    )
+
+    obs = {
+        f"{OBS_STR}.policy": {
+            "present_key": torch.randn(BATCH_SIZE, 10),
+            # missing_key not present
+        },
+    }
+
+    processed = processor.observation(obs)
+
+    # Should only have present_key
+    assert processed[OBS_STATE].shape == (BATCH_SIZE, 10)
+
+
+# =============================================================================
+# Camera/Image Processing Tests
+# =============================================================================
+
+
+def test_camera_permutation_bhwc_to_bchw(processor, sample_observation):
+    """Test images are permuted from (B, H, W, C) to (B, C, H, W)."""
+    processed = processor.observation(sample_observation)
+
+    img_key = f"{OBS_IMAGES}.{TEST_CAMERA_KEY}"
+    assert img_key in processed
+    img = processed[img_key]
+    assert img.shape == (BATCH_SIZE, 3, IMG_HEIGHT, IMG_WIDTH)
+
+
+def test_camera_uint8_to_normalized_float32(processor):
+    """Test that uint8 images are normalized to float32 [0, 1]."""
+    obs = {
+        f"{OBS_STR}.camera_obs": {
+            TEST_CAMERA_KEY: torch.full((BATCH_SIZE, IMG_HEIGHT, IMG_WIDTH, 3), 255, dtype=torch.uint8),
+        },
+    }
+
+    processed = processor.observation(obs)
+
+    img = processed[f"{OBS_IMAGES}.{TEST_CAMERA_KEY}"]
+    assert img.dtype == torch.float32
+    assert torch.allclose(img, torch.ones_like(img))
+
+
+def test_camera_float32_passthrough(processor):
+    """Test that float32 images are kept as float32."""
+    original_img = torch.rand(BATCH_SIZE, IMG_HEIGHT, IMG_WIDTH, 3, dtype=torch.float32)
+    obs = {
+        f"{OBS_STR}.camera_obs": {
+            TEST_CAMERA_KEY: original_img.clone(),
+        },
+    }
+
+    processed = processor.observation(obs)
+
+    img = processed[f"{OBS_IMAGES}.{TEST_CAMERA_KEY}"]
+    assert img.dtype == torch.float32
+    # Values should be same (just permuted)
+    expected = original_img.permute(0, 3, 1, 2)
+    assert torch.allclose(img, expected)
+
+
+def test_camera_other_dtype_converted_to_float(processor):
+    """Test that non-uint8, non-float32 dtypes are converted to float."""
+    obs = {
+        f"{OBS_STR}.camera_obs": {
+            TEST_CAMERA_KEY: torch.randint(0, 255, (BATCH_SIZE, IMG_HEIGHT, IMG_WIDTH, 3), dtype=torch.int32),
+        },
+    }
+
+    processed = processor.observation(obs)
+
+    img = processed[f"{OBS_IMAGES}.{TEST_CAMERA_KEY}"]
+    assert img.dtype == torch.float32
+
+
+def test_camera_filters_to_configured_keys():
+    """Test that only configured camera_keys are extracted."""
+    processor = IsaaclabArenaProcessorStep(
+        state_keys=(),
+        camera_keys=("included_cam",),
+    )
+
+    obs = {
+        f"{OBS_STR}.camera_obs": {
+            "included_cam": torch.randint(0, 255, (BATCH_SIZE, IMG_HEIGHT, IMG_WIDTH, 3), dtype=torch.uint8),
+            "excluded_cam": torch.randint(0, 255, (BATCH_SIZE, IMG_HEIGHT, IMG_WIDTH, 3), dtype=torch.uint8),
+        },
+    }
+
+    processed = processor.observation(obs)
+
+    assert f"{OBS_IMAGES}.included_cam" in processed
+    assert f"{OBS_IMAGES}.excluded_cam" not in processed
+
+
+def test_camera_key_preserved_exactly():
+    """Test that camera key name is used exactly (no suffix stripping)."""
+    processor = IsaaclabArenaProcessorStep(
+        state_keys=(),
+        camera_keys=("my_cam_rgb",),
+    )
+
+    obs = {
+        f"{OBS_STR}.camera_obs": {
+            "my_cam_rgb": torch.randint(0, 255, (BATCH_SIZE, IMG_HEIGHT, IMG_WIDTH, 3), dtype=torch.uint8),
+        },
+    }
+
+    processed = processor.observation(obs)
+
+    # Key should be exactly as configured, with _rgb suffix intact
+    assert f"{OBS_IMAGES}.my_cam_rgb" in processed
+    assert f"{OBS_IMAGES}.my_cam" not in processed
+
+
+# =============================================================================
+# Edge Cases & Missing Data Tests
+# =============================================================================
+
+
+def test_missing_camera_obs_section(processor):
+    """Test processor handles observation without camera_obs section."""
+    obs = {
+        f"{OBS_STR}.policy": {
+            TEST_STATE_KEY: torch.randn(BATCH_SIZE, STATE_DIM),
+        },
+    }
+
+    processed = processor.observation(obs)
+
+    assert OBS_STATE in processed
+    assert not any(k.startswith(OBS_IMAGES) for k in processed)
+
+
+def test_missing_policy_obs_section(processor):
+    """Test processor handles observation without policy section."""
+    obs = {
+        f"{OBS_STR}.camera_obs": {
+            TEST_CAMERA_KEY: torch.randint(0, 255, (BATCH_SIZE, IMG_HEIGHT, IMG_WIDTH, 3), dtype=torch.uint8),
+        },
+    }
+
+    processed = processor.observation(obs)
+
+    assert f"{OBS_IMAGES}.{TEST_CAMERA_KEY}" in processed
+    assert OBS_STATE not in processed
+
+
+def test_empty_observation(processor):
+    """Test processor handles empty observation dict."""
+    processed = processor.observation({})
+
+    assert OBS_STATE not in processed
+    assert not any(k.startswith(OBS_IMAGES) for k in processed)
+
+
+def test_no_matching_state_keys():
+    """Test processor when no state keys match observation."""
+    processor = IsaaclabArenaProcessorStep(
+        state_keys=("nonexistent_key",),
+        camera_keys=(),
+    )
+
+    obs = {
+        f"{OBS_STR}.policy": {
+            "some_other_key": torch.randn(BATCH_SIZE, STATE_DIM),
+        },
+    }
+
+    processed = processor.observation(obs)
+
+    # No state because no keys matched
+    assert OBS_STATE not in processed
+
+
+def test_no_matching_camera_keys():
+    """Test processor when no camera keys match observation."""
+    processor = IsaaclabArenaProcessorStep(
+        state_keys=(),
+        camera_keys=("nonexistent_cam",),
+    )
+
+    obs = {
+        f"{OBS_STR}.camera_obs": {
+            "some_other_cam": torch.randint(
+                0, 255, (BATCH_SIZE, IMG_HEIGHT, IMG_WIDTH, 3), dtype=torch.uint8
+            ),
+        },
+    }
+
+    processed = processor.observation(obs)
+
+    assert not any(k.startswith(OBS_IMAGES) for k in processed)
+
+
+# =============================================================================
+# Configuration Tests
+# =============================================================================
+
+
+def test_default_keys():
+    """Test default state_keys and camera_keys values."""
+    processor = IsaaclabArenaProcessorStep()
+
+    assert processor.state_keys == ("robot_joint_pos",)
+    assert processor.camera_keys == ("robot_pov_cam_rgb",)
+
+
+def test_custom_keys_configuration():
+    """Test processor with custom state and camera keys."""
+    processor = IsaaclabArenaProcessorStep(
+        state_keys=("pos_xyz", "quat_wxyz", "grip_val"),
+        camera_keys=("front_view", "wrist_view"),
+    )
+
+    obs = {
+        f"{OBS_STR}.policy": {
+            "pos_xyz": torch.randn(BATCH_SIZE, 3),
+            "quat_wxyz": torch.randn(BATCH_SIZE, 4),
+            "grip_val": torch.randn(BATCH_SIZE, 1),
+        },
+        f"{OBS_STR}.camera_obs": {
+            "front_view": torch.randint(0, 255, (BATCH_SIZE, IMG_HEIGHT, IMG_WIDTH, 3), dtype=torch.uint8),
+            "wrist_view": torch.randint(0, 255, (BATCH_SIZE, IMG_HEIGHT, IMG_WIDTH, 3), dtype=torch.uint8),
+        },
+    }
+
+    processed = processor.observation(obs)
+
+    # State should be concatenated: 3 + 4 + 1 = 8
+    assert processed[OBS_STATE].shape == (BATCH_SIZE, 8)
+    # Both cameras should be present
+    assert f"{OBS_IMAGES}.front_view" in processed
+    assert f"{OBS_IMAGES}.wrist_view" in processed
+
+
+# =============================================================================
+# transform_features Tests
+# =============================================================================
+
+
+def test_transform_features_passthrough(processor):
+    """Test that transform_features returns features unchanged."""
+    input_features = {
+        PipelineFeatureType.OBSERVATION: {
+            "observation.state": PolicyFeature(
+                type=FeatureType.STATE,
+                shape=(16,),
+            ),
+            "observation.images.cam": PolicyFeature(
+                type=FeatureType.VISUAL,
+                shape=(3, 64, 64),
+            ),
+        },
+        PipelineFeatureType.ACTION: {
+            "action": PolicyFeature(type=FeatureType.ACTION, shape=(7,)),
+        },
+    }
+
+    output_features = processor.transform_features(input_features)
+
+    # Should be unchanged
+    assert output_features == input_features

tests/test_envs.py

@@ -0,0 +1,151 @@
+from isaaclab_env_wrapper import IsaacLabEnvWrapper
+from unittest.mock import MagicMock
+import gym
+import numpy as np
+import torch
+
+
+def _create_mock_isaaclab_env(num_envs: int = 2, device: str = "cpu"):
+    """Create a mock IsaacLab environment for testing."""
+    mock_env = MagicMock()
+    mock_env.num_envs = num_envs
+    mock_env.device = device
+    mock_env.observation_space = gym.spaces.Dict(
+        {"policy": gym.spaces.Box(low=-1, high=1, shape=(num_envs, 54), dtype=np.float32)}
+    )
+    mock_env.action_space = gym.spaces.Box(low=-1, high=1, shape=(36,), dtype=np.float32)
+    mock_env.metadata = {}
+    return mock_env
+
+
+def test_isaaclab_wrapper_init():
+    """Test IsaacLabEnvWrapper initialization."""
+    mock_env = _create_mock_isaaclab_env(num_envs=4)
+
+    wrapper = IsaacLabEnvWrapper(
+        mock_env,
+        episode_length=300,
+        task="Test task",
+        render_mode="rgb_array",
+    )
+
+    assert wrapper.num_envs == 4
+    assert wrapper._max_episode_steps == 300
+    assert wrapper.task == "Test task"
+    assert wrapper.render_mode == "rgb_array"
+    assert wrapper.device == "cpu"
+    assert len(wrapper.envs) == 4
+
+
+def test_isaaclab_wrapper_reset():
+    """Test IsaacLabEnvWrapper reset."""
+    mock_env = _create_mock_isaaclab_env(num_envs=2)
+    mock_obs = {"policy": torch.randn(2, 54)}
+    mock_env.reset.return_value = (mock_obs, {})
+
+    wrapper = IsaacLabEnvWrapper(mock_env, episode_length=100)
+    obs, info = wrapper.reset(seed=42)
+
+    mock_env.reset.assert_called_once_with(seed=42, options=None)
+    assert "final_info" in info
+    assert "is_success" in info["final_info"]
+    assert len(info["final_info"]["is_success"]) == 2
+
+
+def test_isaaclab_wrapper_reset_with_seed_list():
+    """Test that seed list is handled correctly (IsaacLab expects single seed)."""
+    mock_env = _create_mock_isaaclab_env(num_envs=2)
+    mock_env.reset.return_value = ({"policy": torch.randn(2, 54)}, {})
+
+    wrapper = IsaacLabEnvWrapper(mock_env)
+    wrapper.reset(seed=[42, 43, 44])
+
+    # Should extract first seed
+    mock_env.reset.assert_called_once_with(seed=42, options=None)
+
+
+def test_isaaclab_wrapper_step():
+    """Test IsaacLabEnvWrapper step."""
+    mock_env = _create_mock_isaaclab_env(num_envs=2)
+    mock_env.step.return_value = (
+        {"policy": torch.randn(2, 54)},
+        torch.tensor([0.5, 0.3]),
+        torch.tensor([False, False]),
+        torch.tensor([False, True]),
+        {},
+    )
+    # Mock termination manager
+    mock_env.termination_manager.get_term.return_value = torch.tensor([False, True])
+
+    wrapper = IsaacLabEnvWrapper(mock_env)
+    actions = np.random.randn(2, 36).astype(np.float32)
+    obs, reward, terminated, truncated, info = wrapper.step(actions)
+
+    assert reward.dtype == np.float32
+    assert terminated.dtype == bool
+    assert truncated.dtype == bool
+    assert len(reward) == 2
+    assert "final_info" in info
+    assert "is_success" in info["final_info"]
+
+
+def test_isaaclab_wrapper_call_method():
+    """Test IsaacLabEnvWrapper call method."""
+    mock_env = _create_mock_isaaclab_env(num_envs=3)
+
+    wrapper = IsaacLabEnvWrapper(mock_env, episode_length=200, task="My task")
+
+    # Test _max_episode_steps
+    result = wrapper.call("_max_episode_steps")
+    assert result == [200, 200, 200]
+
+    # Test task
+    result = wrapper.call("task")
+    assert result == ["My task", "My task", "My task"]
+
+
+def test_isaaclab_wrapper_render():
+    """Test IsaacLabEnvWrapper render."""
+    mock_env = _create_mock_isaaclab_env(num_envs=2)
+    mock_frames = torch.randint(0, 255, (2, 480, 640, 3), dtype=torch.uint8)
+    mock_env.render.return_value = mock_frames
+
+    wrapper = IsaacLabEnvWrapper(mock_env, render_mode="rgb_array")
+    frame = wrapper.render()
+
+    assert frame is not None
+    assert frame.shape == (480, 640, 3)  # Returns first env frame
+
+
+def test_isaaclab_wrapper_render_all():
+    """Test IsaacLabEnvWrapper render_all."""
+    mock_env = _create_mock_isaaclab_env(num_envs=2)
+    mock_frames = torch.randint(0, 255, (2, 480, 640, 3), dtype=torch.uint8)
+    mock_env.render.return_value = mock_frames
+
+    wrapper = IsaacLabEnvWrapper(mock_env, render_mode="rgb_array")
+    frames = wrapper.render_all()
+
+    assert len(frames) == 2
+    assert all(f.shape == (480, 640, 3) for f in frames)
+
+
+def test_isaaclab_wrapper_render_none():
+    """Test render returns None when render_mode is not rgb_array."""
+    mock_env = _create_mock_isaaclab_env()
+
+    wrapper = IsaacLabEnvWrapper(mock_env, render_mode=None)
+    assert wrapper.render() is None
+
+
+def test_isaaclab_wrapper_close():
+    """Test IsaacLabEnvWrapper close."""
+    mock_env = _create_mock_isaaclab_env()
+    mock_app = MagicMock()
+
+    wrapper = IsaacLabEnvWrapper(mock_env, simulation_app=mock_app)
+    wrapper.close()
+
+    mock_env.close.assert_called_once()
+    mock_app.app.close.assert_called_once()
+    assert wrapper._closed