import json import os import tempfile from unittest.mock import Mock, patch import pytest import yaml from vla_foundry.data.processor.robotics_processor import RoboticsProcessor from vla_foundry.data.robotics.normalization import RoboticsNormalizer from vla_foundry.params.data_params import RoboticsDataParams from vla_foundry.params.robotics.normalization_params import NormalizationParams @pytest.fixture def dataset_stats_path(): """Get the path to the real dataset statistics file.""" return os.path.join(os.path.dirname(__file__), "..", "test_assets", "small_lbm_dataset", "stats.json") @pytest.fixture def dataset_manifest_path(): """Get the path to the dataset manifest file.""" return os.path.join(os.path.dirname(__file__), "..", "test_assets", "small_lbm_dataset", "manifest.jsonl") class TestRoboticsProcessorLoad: """Test the load() and from_pretrained() methods of RoboticsProcessor.""" @staticmethod def _build_normalization_config(action_fields, proprio_fields, **overrides): method = overrides.get("method", "std") scope = overrides.get("scope", "global") epsilon = overrides.get("epsilon", 1e-8) enabled = overrides.get("enabled", True) lowdim_past = overrides.get("lowdim_past_timesteps", 1) lowdim_future = overrides.get("lowdim_future_timesteps", 14) # Deduplicate fields while preserving order (Python 3.7+ idiom) include_fields = list(dict.fromkeys(proprio_fields + action_fields)) field_configs = { field: { "method": method, "scope": scope, "epsilon": epsilon, "enabled": enabled, } for field in include_fields } return { "enabled": enabled, "method": method, "scope": scope, "epsilon": epsilon, "include_fields": include_fields, "field_configs": field_configs, "lowdim_past_timesteps": lowdim_past, "lowdim_future_timesteps": lowdim_future, } @pytest.fixture def sample_config_data(self, dataset_stats_path, dataset_manifest_path): """Create sample configuration data for testing.""" proprio_fields = [ "robot__actual__joint_position__right::panda", "robot__actual__joint_velocity__right::panda", ] action_fields = ["robot__action__poses__right::panda__xyz"] return { "dataset_statistics": [dataset_stats_path], "dataset_manifest": [dataset_manifest_path], "processor": "google/paligemma-3b-pt-224", "proprioception_fields": proprio_fields, "action_fields": action_fields, "normalization": self._build_normalization_config(action_fields, proprio_fields), } @pytest.fixture def sample_statistics_data(self, dataset_stats_path): """Create sample statistics data for testing using real dataset structure.""" # Load the full statistics file with open(dataset_stats_path) as f: full_stats = json.load(f) # Return a subset of the most relevant fields for testing return { "robot__actual__joint_position__right::panda": full_stats["robot__actual__joint_position__right::panda"], "robot__actual__joint_velocity__right::panda": full_stats["robot__actual__joint_velocity__right::panda"], "robot__action__poses__right::panda__xyz": full_stats["robot__action__poses__right::panda__xyz"], } @pytest.fixture def temp_config_file(self, dataset_stats_path, dataset_manifest_path): """Create a temporary config file for testing.""" proprio_fields = [ "robot__actual__joint_position__right::panda", "robot__actual__joint_velocity__right::panda", ] action_fields = ["robot__actual__poses__right::panda__xyz"] config_data = { "dataset_statistics": [dataset_stats_path], "dataset_manifest": [dataset_manifest_path], "processor": "google/paligemma-3b-pt-224", "proprioception_fields": proprio_fields, "action_fields": action_fields, "normalization": self._build_normalization_config(action_fields, proprio_fields), "lowdim_past_timesteps": 1, "lowdim_future_timesteps": 14, } with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: yaml.safe_dump(config_data, f) temp_path = f.name yield temp_path # Cleanup if os.path.exists(temp_path): os.unlink(temp_path) @pytest.fixture def temp_stats_file(self, sample_statistics_data): """Create a temporary statistics file for testing.""" with tempfile.TemporaryDirectory() as temp_dir: stats_path = os.path.join(temp_dir, "stats.json") with open(stats_path, "w") as f: json.dump(sample_statistics_data, f) metadata_path = os.path.join(temp_dir, "processing_metadata.json") with open(metadata_path, "w") as f: json.dump({"processing": {"past_lowdim_steps": 1, "future_lowdim_steps": 14}}, f) yield stats_path @pytest.fixture def temp_experiment_dir(self, dataset_stats_path, dataset_manifest_path): """Create a temporary experiment directory with config and stats files.""" with tempfile.TemporaryDirectory() as temp_dir: proprio_fields = [ "robot__actual__joint_position__right::panda", "robot__actual__joint_velocity__right::panda", ] action_fields = ["robot__actual__poses__right::panda__xyz"] normalization_cfg = self._build_normalization_config(action_fields, proprio_fields) # Create config_processor.yaml config_path = os.path.join(temp_dir, "config_processor.yaml") config_data = { "lowdim_past_timesteps": 0, "lowdim_future_timesteps": 8, "dataset_statistics": [dataset_stats_path], "dataset_manifest": [dataset_manifest_path], "processor": "google/paligemma-3b-pt-224", "proprioception_fields": proprio_fields, "action_fields": action_fields, "normalization": normalization_cfg, } with open(config_path, "w") as f: yaml.safe_dump(config_data, f) # Create config_normalizer.yaml and stats.json required by RoboticsNormalizer.from_pretrained normalizer_config_path = os.path.join(temp_dir, "config_normalizer.yaml") with open(normalizer_config_path, "w") as f: yaml.safe_dump(normalization_cfg, f) with open(dataset_stats_path) as f: stats_data = json.load(f) with open(os.path.join(temp_dir, "stats.json"), "w") as f: json.dump(stats_data, f) yield temp_dir @patch("vla_foundry.data.processor.robotics_processor.get_processor") def test_robotics_processor_load(self, mock_get_processor, temp_config_file): """Test RoboticsProcessor.load() method.""" # Setup mocks mock_processor = Mock() mock_get_processor.return_value = mock_processor # Test load method (now using real dataset statistics file) processor = RoboticsProcessor.load(temp_config_file) # Assertions assert isinstance(processor, RoboticsProcessor) assert isinstance(processor.data_params, RoboticsDataParams) assert processor.data_params.type == "robotics" assert processor.data_params.processor == "google/paligemma-3b-pt-224" assert processor.data_params.proprioception_fields == [ "robot__actual__joint_position__right::panda", "robot__actual__joint_velocity__right::panda", ] assert processor.data_params.action_fields == ["robot__actual__poses__right::panda__xyz"] assert processor.data_params.normalization.enabled is True # Verify processor was initialized mock_get_processor.assert_called_once() # Verify normalizer was created since normalization is enabled assert processor.normalizer is not None assert isinstance(processor.normalizer, RoboticsNormalizer) assert processor.normalizer.lowdim_past_timesteps == 1 assert processor.normalizer.lowdim_future_timesteps == 14 assert processor.data_params.normalization.lowdim_past_timesteps == 1 assert processor.data_params.normalization.lowdim_future_timesteps == 14 assert processor.data_params.lowdim_past_timesteps == 1 assert processor.data_params.lowdim_future_timesteps == 14 # Verify that real statistics were loaded assert processor.normalizer.stats is not None expected_fields = [ "robot__actual__joint_position__right::panda", "robot__actual__joint_velocity__right::panda", "robot__actual__poses__right::panda__xyz", ] for field in expected_fields: assert field in processor.normalizer.stats, f"Field {field} not found in statistics" @patch("vla_foundry.data.processor.robotics_processor.get_processor") def test_robotics_processor_from_pretrained(self, mock_get_processor, temp_experiment_dir): """Test RoboticsProcessor.from_pretrained() method.""" # Setup mocks mock_processor = Mock() mock_get_processor.return_value = mock_processor # Test from_pretrained method using dataset statistics file processor = RoboticsProcessor.from_pretrained(temp_experiment_dir) # Assertions assert isinstance(processor, RoboticsProcessor) assert isinstance(processor.data_params, RoboticsDataParams) assert processor.data_params.type == "robotics" assert processor.data_params.processor == "google/paligemma-3b-pt-224" # Verify processor was initialized mock_get_processor.assert_called_once() # Verify normalizer was created and has real statistics assert processor.normalizer is not None assert isinstance(processor.normalizer, RoboticsNormalizer) assert processor.normalizer.stats is not None assert processor.data_params.lowdim_past_timesteps == 0 assert processor.data_params.lowdim_future_timesteps == 8 assert processor.data_params.normalization.lowdim_past_timesteps == 1 assert processor.data_params.normalization.lowdim_future_timesteps == 14 assert processor.normalizer.lowdim_past_timesteps == 1 assert processor.normalizer.lowdim_future_timesteps == 14 # Verify that real statistics were loaded expected_fields = [ "robot__actual__joint_position__right::panda", "robot__actual__joint_velocity__right::panda", "robot__actual__poses__right::panda__xyz", ] for field in expected_fields: assert field in processor.normalizer.stats, f"Field {field} not found in statistics" @patch("vla_foundry.data.processor.robotics_processor.get_processor") def test_robotics_processor_load_with_normalization_disabled( self, mock_get_processor, dataset_stats_path, dataset_manifest_path ): """Test RoboticsProcessor.load() with normalization disabled.""" # Setup mocks mock_processor = Mock() mock_get_processor.return_value = mock_processor # Create config with normalization disabled with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: f.write("dataset_statistics:\n") f.write(f" - {dataset_stats_path}\n") f.write("dataset_manifest:\n") f.write(f" - {dataset_manifest_path}\n") f.write("processor: google/paligemma-3b-pt-224\n") f.write("proprioception_fields: []\n") f.write("action_fields: []\n") f.write("normalization:\n") f.write(" enabled: false\n") temp_config_path = f.name try: # Test load method processor = RoboticsProcessor.load(temp_config_path) # Assertions assert isinstance(processor, RoboticsProcessor) assert processor.data_params.normalization.enabled is False assert processor.normalizer is None finally: # Cleanup temp config file if os.path.exists(temp_config_path): os.unlink(temp_config_path) def test_robotics_processor_load_nonexistent_file(self): """Test RoboticsProcessor.load() with nonexistent config file.""" # draccus.load throws a different exception for nonexistent files with pytest.raises((FileNotFoundError, Exception)): RoboticsProcessor.load("/nonexistent/config.yaml") def test_robotics_processor_from_pretrained_nonexistent_dir(self): """Test RoboticsProcessor.from_pretrained() with nonexistent directory.""" # draccus.load throws a different exception for nonexistent files with pytest.raises((FileNotFoundError, Exception)): RoboticsProcessor.from_pretrained("/nonexistent/dir") @patch("vla_foundry.data.processor.robotics_processor.get_processor") def test_robotics_processor_with_actual_dataset_statistics( self, mock_get_processor, dataset_stats_path, dataset_manifest_path ): """Test RoboticsProcessor with actual dataset statistics structure.""" # Setup mock mock_processor = Mock() mock_get_processor.return_value = mock_processor with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: proprio_fields = [ "robot__actual__joint_position__right::panda", "robot__actual__joint_velocity__right::panda", ] action_fields = ["robot__actual__poses__right::panda__xyz"] normalization_cfg = self._build_normalization_config(action_fields, proprio_fields) normalization_cfg["field_configs"]["robot__actual__joint_position__right::panda"] = { "method": "percentile_5_95", "scope": "per_timestep", "epsilon": 1e-6, "enabled": True, } config_data = { "dataset_statistics": [dataset_stats_path], "dataset_manifest": [dataset_manifest_path], "processor": "google/paligemma-3b-pt-224", "proprioception_fields": proprio_fields, "action_fields": action_fields, "normalization": normalization_cfg, } yaml.safe_dump(config_data, f) temp_config_path = f.name try: # Test load method with actual dataset statistics processor = RoboticsProcessor.load(temp_config_path) # Assertions assert isinstance(processor, RoboticsProcessor) assert isinstance(processor.data_params, RoboticsDataParams) assert processor.data_params.type == "robotics" # Verify the processor was initialized mock_get_processor.assert_called_once() # Verify normalizer was created and has real statistics assert processor.normalizer is not None assert isinstance(processor.normalizer, RoboticsNormalizer) assert processor.normalizer.stats is not None # Check that the actual dataset fields are present in statistics expected_fields = [ "robot__actual__joint_position__right::panda", "robot__actual__joint_velocity__right::panda", "robot__actual__poses__right::panda__xyz", ] for field in expected_fields: assert field in processor.normalizer.stats, f"Field {field} not found in statistics" # Check that each field has the expected statistics structure field_stats = processor.normalizer.stats[field] assert "mean" in field_stats assert "std" in field_stats assert "min" in field_stats assert "max" in field_stats assert "mean_per_timestep" in field_stats assert "std_per_timestep" in field_stats # Verify the statistics are lists/arrays (not scalars) assert isinstance(field_stats["mean"], list) assert isinstance(field_stats["std"], list) assert isinstance(field_stats["min"], list) assert isinstance(field_stats["max"], list) assert isinstance(field_stats["mean_per_timestep"], list) assert isinstance(field_stats["std_per_timestep"], list) # Verify the per-timestep data has the expected structure assert len(field_stats["mean_per_timestep"]) > 0, f"No timestep data for {field}" assert len(field_stats["std_per_timestep"]) > 0, f"No timestep std data for {field}" finally: # Cleanup if os.path.exists(temp_config_path): os.unlink(temp_config_path) class TestRoboticsNormalizerLoad: """Test the load() and from_pretrained() methods of RoboticsNormalizer.""" @pytest.fixture def sample_normalization_config_data(self): """Create sample normalization configuration data for testing.""" fields = ["robot__actual__joint_position__right::panda"] return { "enabled": True, "method": "std", "scope": "global", "epsilon": 1e-8, "include_fields": fields, "lowdim_past_timesteps": 1, "lowdim_future_timesteps": 14, "field_configs": { "robot__actual__joint_position__right::panda": { "method": "percentile_5_95", "scope": "per_timestep", "epsilon": 1e-6, } }, } @pytest.fixture def sample_statistics_data(self): """Create sample statistics data for testing using real dataset structure.""" # Load a subset of the real dataset statistics for testing dataset_stats_path = os.path.join( os.path.dirname(__file__), "..", "test_assets", "small_lbm_dataset", "stats.json" ) # Load the full statistics file with open(dataset_stats_path) as f: full_stats = json.load(f) # Return a subset of the most relevant fields for testing return { "robot__actual__joint_position__right::panda": full_stats["robot__actual__joint_position__right::panda"], "robot__actual__joint_velocity__right::panda": full_stats["robot__actual__joint_velocity__right::panda"], "robot__actual__poses__right::panda__xyz": full_stats["robot__actual__poses__right::panda__xyz"], } @pytest.fixture def temp_normalization_config_file(self): """Create a temporary normalization config file for testing.""" with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: f.write("enabled: true\n") f.write("method: std\n") f.write("scope: global\n") f.write("epsilon: 1.0e-08\n") f.write("lowdim_past_timesteps: 1\n") f.write("lowdim_future_timesteps: 14\n") f.write("field_configs:\n") f.write(" robot__actual__joint_position__right::panda:\n") f.write(" method: percentile_5_95\n") f.write(" scope: per_timestep\n") f.write(" epsilon: 1.0e-06\n") temp_path = f.name yield temp_path # Cleanup if os.path.exists(temp_path): os.unlink(temp_path) @pytest.fixture def temp_stats_file(self, sample_statistics_data): """Create a temporary statistics file for testing.""" with tempfile.TemporaryDirectory() as temp_dir: stats_path = os.path.join(temp_dir, "stats.json") with open(stats_path, "w") as f: json.dump(sample_statistics_data, f) metadata_path = os.path.join(temp_dir, "processing_metadata.json") with open(metadata_path, "w") as f: json.dump({"processing": {"past_lowdim_steps": 1, "future_lowdim_steps": 14}}, f) yield stats_path @pytest.fixture def temp_experiment_dir(self, sample_statistics_data): """Create a temporary experiment directory with normalizer config and stats files.""" with tempfile.TemporaryDirectory() as temp_dir: # Create config_normalizer.yaml config_path = os.path.join(temp_dir, "config_normalizer.yaml") with open(config_path, "w") as f: f.write("enabled: true\n") f.write("method: std\n") f.write("scope: global\n") f.write("epsilon: 1.0e-08\n") f.write("lowdim_past_timesteps: 1\n") f.write("lowdim_future_timesteps: 14\n") f.write("field_configs:\n") f.write(" robot__actual__joint_position__right::panda:\n") f.write(" method: percentile_5_95\n") f.write(" scope: per_timestep\n") f.write(" epsilon: 1.0e-06\n") # Create stats.json stats_path = os.path.join(temp_dir, "stats.json") with open(stats_path, "w") as f: json.dump(sample_statistics_data, f) metadata_path = os.path.join(temp_dir, "processing_metadata.json") with open(metadata_path, "w") as f: json.dump({"processing": {"past_lowdim_steps": 1, "future_lowdim_steps": 14}}, f) yield temp_dir def test_robotics_normalizer_load(self, temp_normalization_config_file, temp_stats_file, dataset_stats_path): """Test RoboticsNormalizer.load() method.""" # Test that the load method works correctly normalizer = RoboticsNormalizer.load(temp_normalization_config_file, temp_stats_file) # Assertions assert isinstance(normalizer, RoboticsNormalizer) assert isinstance(normalizer.normalization_params, NormalizationParams) assert normalizer.normalization_params.enabled is True assert normalizer.normalization_params.method == "std" assert normalizer.normalization_params.scope == "global" assert normalizer.normalization_params.epsilon == 1e-8 # Check field-specific config assert "robot__actual__joint_position__right::panda" in normalizer.normalization_params.field_configs joint_pos_config = normalizer.normalization_params.field_configs["robot__actual__joint_position__right::panda"] assert joint_pos_config.method == "percentile_5_95" assert joint_pos_config.scope == "per_timestep" assert joint_pos_config.epsilon == 1e-6 # Verify statistics were loaded assert normalizer.stats is not None assert "robot__actual__joint_position__right::panda" in normalizer.stats assert "robot__actual__joint_velocity__right::panda" in normalizer.stats assert normalizer.lowdim_past_timesteps == 1 assert normalizer.lowdim_future_timesteps == 14 assert normalizer.normalization_params.lowdim_past_timesteps == 1 assert normalizer.normalization_params.lowdim_future_timesteps == 14 assert normalizer.lowdim_past_timesteps == 1 assert normalizer.lowdim_future_timesteps == 14 assert normalizer.normalization_params.lowdim_past_timesteps == 1 assert normalizer.normalization_params.lowdim_future_timesteps == 14 def test_robotics_normalizer_from_pretrained(self, temp_experiment_dir, dataset_stats_path): """Test RoboticsNormalizer.from_pretrained() method.""" # Test that the from_pretrained method works correctly normalizer = RoboticsNormalizer.from_pretrained(temp_experiment_dir) # Assertions assert isinstance(normalizer, RoboticsNormalizer) assert isinstance(normalizer.normalization_params, NormalizationParams) assert normalizer.normalization_params.enabled is True assert normalizer.normalization_params.method == "std" assert normalizer.normalization_params.scope == "global" assert normalizer.normalization_params.epsilon == 1e-8 # Check field-specific config assert "robot__actual__joint_position__right::panda" in normalizer.normalization_params.field_configs joint_pos_config = normalizer.normalization_params.field_configs["robot__actual__joint_position__right::panda"] assert joint_pos_config.method == "percentile_5_95" assert joint_pos_config.scope == "per_timestep" assert joint_pos_config.epsilon == 1e-6 # Verify statistics were loaded assert normalizer.stats is not None assert "robot__actual__joint_position__right::panda" in normalizer.stats assert "robot__actual__joint_velocity__right::panda" in normalizer.stats def test_robotics_normalizer_load_with_disabled_normalization(self, temp_stats_file, dataset_stats_path): """Test RoboticsNormalizer.load() with disabled normalization.""" # Create normalization config with normalization disabled with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: f.write("enabled: false\n") f.write("method: std\n") f.write("scope: global\n") f.write("epsilon: 1.0e-08\n") f.write("lowdim_past_timesteps: 1\n") f.write("lowdim_future_timesteps: 14\n") f.write("field_configs: {}\n") temp_config_path = f.name try: # Test that load works with disabled normalization normalizer = RoboticsNormalizer.load(temp_config_path, temp_stats_file) # Assertions assert isinstance(normalizer, RoboticsNormalizer) assert normalizer.normalization_params.enabled is False assert normalizer.stats is not None # even with disabled normalization, stats are loaded assert normalizer.enabled is False assert normalizer.lowdim_past_timesteps == 1 assert normalizer.lowdim_future_timesteps == 14 finally: # Cleanup if os.path.exists(temp_config_path): os.unlink(temp_config_path) def test_robotics_normalizer_load_nonexistent_config(self, temp_stats_file): """Test RoboticsNormalizer.load() with nonexistent config file.""" # The actual load method has a bug, but we test the error it would produce # when trying to load a nonexistent file with pytest.raises((FileNotFoundError, AttributeError)): RoboticsNormalizer.load("/nonexistent/config.yaml", temp_stats_file) def test_robotics_normalizer_load_nonexistent_stats(self, temp_normalization_config_file): """Test RoboticsNormalizer.load() with nonexistent statistics file.""" # The load method has a bug, but we can still test that it would fail with pytest.raises((FileNotFoundError, AttributeError)): RoboticsNormalizer.load(temp_normalization_config_file, "/nonexistent/stats.json") def test_robotics_normalizer_from_pretrained_nonexistent_dir(self): """Test RoboticsNormalizer.from_pretrained() with nonexistent directory.""" # The from_pretrained method has a bug, but we can still test error handling with pytest.raises((FileNotFoundError, AttributeError)): RoboticsNormalizer.from_pretrained("/nonexistent/dir") def test_robotics_normalizer_from_pretrained_missing_config(self, sample_statistics_data): """Test RoboticsNormalizer.from_pretrained() with missing config file.""" with tempfile.TemporaryDirectory() as temp_dir: # Only create stats file, not config file stats_path = os.path.join(temp_dir, "stats.json") with open(stats_path, "w") as f: json.dump(sample_statistics_data, f) # The from_pretrained method has a bug, but we can still test error handling with pytest.raises((FileNotFoundError, AttributeError)): RoboticsNormalizer.from_pretrained(temp_dir) def test_robotics_normalizer_from_pretrained_missing_stats(self): """Test RoboticsNormalizer.from_pretrained() with missing stats file.""" with tempfile.TemporaryDirectory() as temp_dir: # Only create config file, not stats file config_path = os.path.join(temp_dir, "config_normalizer.yaml") with open(config_path, "w") as f: f.write("enabled: true\n") f.write("method: std\n") f.write("scope: global\n") f.write("epsilon: 1.0e-08\n") f.write("field_configs: {}\n") # The from_pretrained method has a bug, but we can still test error handling with pytest.raises((FileNotFoundError, AttributeError)): RoboticsNormalizer.from_pretrained(temp_dir)