init commit

tests/__init__.py

@@ -0,0 +1 @@
+# Tests module for DataEngine

tests/integration/__init__.py

@@ -0,0 +1 @@
+# Integration tests module

tests/integration/base/__init__.py

@@ -0,0 +1 @@
+# Base test infrastructure

tests/integration/base/test_harness.py

@@ -0,0 +1,401 @@
+"""
+Integration Test Harness base class that encapsulates common test logic.
+Reduces boilerplate code and provides a consistent testing interface.
+"""
+
+import os
+import shutil
+import subprocess
+import sys
+from pathlib import Path
+from typing import Any, Dict, Optional, Tuple
+
+from .utils import set_test_seeds
+
+# Add paths to sys.path for proper imports
+sys.path.append("./")
+sys.path.append("./data_engine")
+
+
+# Import data_engine modules only when needed to avoid import errors during framework testing
+def _import_data_engine_modules():
+    """Import data_engine modules when they are actually needed."""
+    try:
+        from nimbus import run_data_engine
+        from nimbus.utils.config_processor import ConfigProcessor
+        from nimbus.utils.utils import init_env
+
+        return init_env, run_data_engine, ConfigProcessor
+    except ImportError as e:
+        raise ImportError(f"Failed to import data_engine modules. Ensure dependencies are installed: {e}")
+
+
+class IntegrationTestHarness:
+    """
+    Base class for integration tests that provides common functionality.
+
+    This class encapsulates:
+    - Configuration loading and processing
+    - Output directory cleanup
+    - Test pipeline execution (direct or subprocess)
+    - Output validation
+    - Data comparison with reference
+    """
+
+    def __init__(self, config_path: str, seed: int = 42, load_num: int = 0, random_num: int = 0, episodes: int = 0):
+        """
+        Initialize the test harness with configuration and seed.
+
+        Args:
+            config_path: Path to the test configuration YAML file
+            seed: Random seed for reproducible results (default: 42)
+        """
+        self.config_path = config_path
+        self.seed = seed
+        self.load_num = load_num
+        self.random_num = random_num
+        self.episodes = episodes
+        self.output_dir = None
+
+        # Import data_engine modules
+        self._init_env, self._run_data_engine, self._ConfigProcessor = _import_data_engine_modules()
+        self.processor = self._ConfigProcessor()
+
+        # Initialize environment (same as launcher.py)
+        self._init_env()
+
+        self.config = self.load_and_process_config() if config_path else None
+        # Set random seeds for reproducibility
+        self.modify_config()
+
+        set_test_seeds(seed)
+
+        from nimbus.utils.flags import set_debug_mode, set_random_seed
+
+        set_debug_mode(True)  # Enable debug mode for better error visibility during tests
+        set_random_seed(seed)
+
+    def modify_config(self):
+        """Modify configuration parameters as needed before running the pipeline."""
+        if self.config and "load_stage" in self.config:
+            if "layout_random_generator" in self.config.load_stage:
+                if "args" in self.config.load_stage.layout_random_generator:
+                    if self.random_num > 0:
+                        self.config.load_stage.layout_random_generator.args.random_num = self.random_num
+                    if "input_dir" in self.config.load_stage.layout_random_generator.args:
+                        if "simbox" in self.config.name:
+                            input_path = (
+                                "/shared/smartbot_new/zhangyuchang/CI/manip/"
+                                "simbox/simbox_plan_ci/seq_path/BananaBaseTask/plan"
+                            )
+                            self.config.load_stage.layout_random_generator.args.input_dir = input_path
+        if self.config and "plan_stage" in self.config:
+            if "seq_planner" in self.config.plan_stage:
+                if "args" in self.config.plan_stage.seq_planner:
+                    if self.episodes > 0:
+                        self.config.plan_stage.seq_planner.args.episodes = self.episodes
+                        if self.load_num > 0:
+                            self.config.load_stage.scene_loader.args.load_num = self.load_num
+        if self.config and "name" in self.config:
+            self.config.name = self.config.name + "_ci"
+        if self.config and "store_stage" in self.config:
+            if hasattr(self.config.store_stage.writer.args, "obs_output_dir"):
+                self.config.store_stage.writer.args.obs_output_dir = f"output/{self.config.name}/obs_path/"
+            if hasattr(self.config.store_stage.writer.args, "seq_output_dir"):
+                self.config.store_stage.writer.args.seq_output_dir = f"output/{self.config.name}/seq_path/"
+            if hasattr(self.config.store_stage.writer.args, "output_dir"):
+                self.config.store_stage.writer.args.output_dir = f"output/{self.config.name}/"
+            self._extract_output_dir()
+
+    def load_and_process_config(self) -> Dict[str, Any]:
+        """
+        Load and process the test configuration file.
+
+        Returns:
+            Processed configuration dictionary
+
+        Raises:
+            AssertionError: If config file does not exist
+        """
+        assert os.path.exists(self.config_path), f"Config file not found: {self.config_path}"
+
+        self.config = self.processor.process_config(self.config_path)
+        self.processor.print_final_config(self.config)
+
+        # Extract output directory from config
+        self._extract_output_dir()
+
+        return self.config
+
+    def _extract_output_dir(self):
+        """Extract and expand the output directory path from config."""
+        # Try navigation test output path
+        output_dir = self.config.get("store_stage", {}).get("writer", {}).get("args", {}).get("seq_output_dir")
+
+        # If not found, try common output_dir used in render configs
+        if not output_dir:
+            output_dir = self.config.get("store_stage", {}).get("writer", {}).get("args", {}).get("output_dir")
+
+        # Process the output directory if found
+        if output_dir and isinstance(output_dir, str):
+            name = self.config.get("name", "test_output")
+            output_dir = output_dir.replace("${name}", name)
+            self.output_dir = os.path.abspath(output_dir)
+
+    def cleanup_output_directory(self):
+        """Clean up existing output directory if it exists."""
+        if self.output_dir and os.path.exists(self.output_dir):
+            # Use ignore_errors=True to handle NFS caching issues
+            shutil.rmtree(self.output_dir, ignore_errors=True)
+            # If directory still exists (NFS delay), try removing with onerror handler
+            if os.path.exists(self.output_dir):
+
+                def handle_remove_error(func, path, exc_info):  # pylint: disable=W0613
+                    """Handle errors during removal, with retry for NFS issues."""
+                    import time
+
+                    time.sleep(0.1)  # Brief delay for NFS sync
+                    try:
+                        if os.path.isdir(path):
+                            os.rmdir(path)
+                        else:
+                            os.remove(path)
+                    except OSError:
+                        pass  # Ignore if still fails
+
+                shutil.rmtree(self.output_dir, onerror=handle_remove_error)
+            print(f"Cleaned up existing output directory: {self.output_dir}")
+
+    def run_data_engine_direct(self, **kwargs) -> None:
+        """
+        Run the test pipeline directly in the current process.
+
+        Args:
+            **kwargs: Additional arguments to pass to run_data_engine
+
+        Raises:
+            Exception: If pipeline execution fails
+        """
+        if not self.config:
+            self.load_and_process_config()
+        self.modify_config()
+        self.cleanup_output_directory()
+
+        try:
+            self._run_data_engine(self.config, **kwargs)
+        except Exception as e:
+            raise e
+
+    def run_data_engine_subprocess(
+        self,
+        runner_script: str,
+        interpreter: str = "python",
+        timeout: int = 1800,
+        compare_output: bool = False,
+        reference_dir: str = "",
+        comparator: str = "simbox",
+        comparator_args: Optional[Dict[str, Any]] = None,
+    ) -> subprocess.CompletedProcess:
+        """
+        Run the test pipeline in a subprocess using a runner script.
+
+        Args:
+            runner_script: Path to the runner script to execute
+            interpreter: Command to use for running the script (default: "python")
+            timeout: Timeout in seconds for subprocess execution (default: 1800)
+            compare_output: Whether to compare generated output with a reference directory
+            reference_dir: Path to reference directory containing meta_info.pkl and lmdb
+            comparator: Which comparator to use (default: "simbox")
+            comparator_args: Optional extra arguments for comparator (e.g. {"tolerance": 1e-6, "image_psnr": 37.0})
+
+        Returns:
+            subprocess.CompletedProcess object with execution results
+
+        Raises:
+            subprocess.TimeoutExpired: If execution exceeds timeout
+            AssertionError: If subprocess returns non-zero exit code
+        """
+        self.cleanup_output_directory()
+
+        # Build command based on interpreter type
+        if interpreter == "blenderproc":
+            cmd = ["blenderproc", "run", runner_script]
+            if os.environ.get("BLENDER_CUSTOM_PATH"):
+                cmd.extend(["--custom-blender-path", os.environ["BLENDER_CUSTOM_PATH"]])
+        elif interpreter.endswith(".sh"):
+            # For scripts like /isaac-sim/python.sh
+            cmd = [interpreter, runner_script]
+        else:
+            cmd = [interpreter, runner_script]
+
+        if not self.output_dir:
+            self._extract_output_dir()
+        output_dir = self.output_dir
+
+        print(f"Running command: {' '.join(cmd)}")
+        print(f"Expected output directory: {output_dir}")
+
+        try:
+            result = subprocess.run(cmd, capture_output=False, text=True, timeout=timeout, check=False)
+
+            # Print subprocess output for debugging
+            if result.stdout:
+                print("STDOUT:", result.stdout[-2000:])  # Last 2000 chars
+            if result.stderr:
+                print("STDERR:", result.stderr[-1000:])  # Last 1000 chars
+            print("Return code:", result.returncode)
+
+            if compare_output and result.returncode == 0:
+                if not output_dir:
+                    raise RuntimeError(
+                        "Output directory not configured. Expected one of "
+                        "store_stage.writer.args.(seq_output_dir|obs_output_dir|output_dir)."
+                    )
+
+                if output_dir:
+                    for root, dirs, files in os.walk(output_dir):
+                        if "lmdb" in dirs and "meta_info.pkl" in files:
+                            output_dir = root
+                            break
+
+                if os.path.exists(reference_dir):
+                    # Find the reference render directory
+                    for root, dirs, files in os.walk(reference_dir):
+                        if "lmdb" in dirs and "meta_info.pkl" in files:
+                            reference_dir = root
+                            break
+
+                # Build comparator command according to requested comparator
+                comp = (comparator or "simbox").lower()
+                comparator_args = comparator_args or {}
+
+                if comp == "simbox":
+                    compare_cmd = [
+                        "/isaac-sim/python.sh",
+                        "tests/integration/data_comparators/simbox_comparator.py",
+                        "--dir1",
+                        output_dir,
+                        "--dir2",
+                        reference_dir,
+                    ]
+                    # Optional numeric/image thresholds
+                    if "tolerance" in comparator_args and comparator_args["tolerance"] is not None:
+                        compare_cmd += ["--tolerance", str(comparator_args["tolerance"])]
+                    if "image_psnr" in comparator_args and comparator_args["image_psnr"] is not None:
+                        compare_cmd += ["--image-psnr", str(comparator_args["image_psnr"])]
+
+                    print(f"Running comparison: {' '.join(compare_cmd)}")
+                    compare_result = subprocess.run(
+                        compare_cmd, capture_output=True, text=True, timeout=600, check=False
+                    )
+
+                    print("Comparison STDOUT:")
+                    print(compare_result.stdout)
+                    print("Comparison STDERR:")
+                    print(compare_result.stderr)
+
+                    if compare_result.returncode != 0:
+                        raise RuntimeError("Simbox comparison failed: outputs differ")
+
+                    if "Successfully loaded data from both directories" in compare_result.stdout:
+                        print("✓ Both output directories have valid structure (meta_info.pkl + lmdb)")
+
+                    print("✓ Simbox render test completed with numeric-aligned comparison")
+
+                else:
+                    raise ValueError(f"Unknown comparator: {comp}. Use 'simbox'.")
+
+            return result
+
+        except subprocess.TimeoutExpired:
+            raise RuntimeError(f"Test timed out after {timeout} seconds")
+
+    def validate_output_generated(self, min_files: int = 1) -> bool:
+        """
+        Validate that output was generated in the expected directory.
+
+        Args:
+            min_files: Minimum number of files expected in output (default: 1)
+
+        Returns:
+            True if validation passes
+
+        Raises:
+            AssertionError: If output directory doesn't exist or has too few files
+        """
+        assert self.output_dir is not None, "Output directory not configured"
+        assert os.path.exists(self.output_dir), f"Expected output directory was not created: {self.output_dir}"
+
+        output_files = list(Path(self.output_dir).rglob("*"))
+        assert (
+            len(output_files) >= min_files
+        ), f"Expected at least {min_files} files but found {len(output_files)} in: {self.output_dir}"
+
+        print(f"✓ Pipeline completed successfully. Output generated in: {self.output_dir}")
+        print(f"✓ Generated {len(output_files)} files/directories")
+        return True
+
+    def compare_with_reference(self, reference_dir: str, comparator_func, **comparator_kwargs) -> Tuple[bool, str]:
+        """
+        Compare generated output with reference data using provided comparator function.
+
+        Args:
+            reference_dir: Path to reference data directory
+            comparator_func: Function to use for comparison
+            **comparator_kwargs: Additional arguments for comparator function
+
+        Returns:
+            Tuple of (success: bool, message: str)
+        """
+        if not os.path.exists(reference_dir):
+            print(f"Reference directory not found, skipping comparison: {reference_dir}")
+            return True, "Reference data not available"
+
+        success, message = comparator_func(self.output_dir, reference_dir, **comparator_kwargs)
+
+        if success:
+            print(f"✓ Results match reference data: {message}")
+        else:
+            print(f"✗ Result comparison failed: {message}")
+
+        return success, message
+
+    def run_test_end_to_end(
+        self,
+        reference_dir: Optional[str] = None,
+        comparator_func=None,
+        comparator_kwargs: Optional[Dict] = None,
+        min_output_files: int = 1,
+        **pipeline_kwargs,
+    ) -> bool:
+        """
+        Run a complete end-to-end test including comparison with reference.
+
+        Args:
+            reference_dir: Optional path to reference data for comparison
+            comparator_func: Optional function for comparing results
+            comparator_kwargs: Optional kwargs for comparator function
+            min_output_files: Minimum expected output files
+            **pipeline_kwargs: Additional arguments for pipeline execution
+
+        Returns:
+            True if test passes, False otherwise
+
+        Raises:
+            Exception: If any test step fails
+        """
+        # Load configuration
+        self.load_and_process_config()
+        # Run pipeline
+        self.run_data_engine_direct(**pipeline_kwargs, master_seed=self.seed)
+
+        # Validate output
+        self.validate_output_generated(min_files=min_output_files)
+
+        # Compare with reference if provided
+        if reference_dir and comparator_func:
+            comparator_kwargs = comparator_kwargs or {}
+            success, message = self.compare_with_reference(reference_dir, comparator_func, **comparator_kwargs)
+            assert success, f"Comparison with reference failed: {message}"
+
+        return True

tests/integration/base/utils.py

@@ -0,0 +1,52 @@
+"""
+Utility functions for integration tests, including centralized seed setting.
+"""
+
+import random
+
+import numpy as np
+
+try:
+    import torch
+except ImportError:
+    torch = None
+
+try:
+    import open3d as o3d
+except ImportError:
+    o3d = None
+
+
+def set_test_seeds(seed):
+    """
+    Set random seeds for all relevant libraries to ensure reproducible results.
+
+    This function sets seeds for:
+    - Python's random module
+    - NumPy
+    - PyTorch (if available)
+    - Open3D (if available)
+    - PyTorch CUDA (if available)
+    - PyTorch CUDNN settings for deterministic behavior
+
+    Args:
+        seed (int): The seed value to use for all random number generators
+    """
+    # Set Python's built-in random seed
+    random.seed(seed)
+
+    # Set NumPy seed
+    np.random.seed(seed)
+
+    # Set PyTorch seeds if available
+    if torch is not None:
+        torch.manual_seed(seed)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(seed)
+        # Configure CUDNN for deterministic behavior
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+
+    # Set Open3D seed if available
+    if o3d is not None:
+        o3d.utility.random.seed(seed)

tests/integration/data_comparators/__init__.py

@@ -0,0 +1,8 @@
+"""
+Data comparators module for integration tests.
+Provides functions to compare generated data with reference data.
+"""
+
+from .sequence_comparator import compare_navigation_results
+
+__all__ = ['compare_navigation_results']

tests/integration/data_comparators/sequence_comparator.py

@@ -0,0 +1,551 @@
+"""
+Sequence data comparator for navigation pipeline tests.
+Provides functions to compare generated navigation sequences with reference data.
+"""
+
+import json
+import os
+import math
+import numpy as np
+import cv2  # OpenCV is available per requirements
+from pathlib import Path
+from typing import Tuple, Any, Dict, Optional
+
+
+def compare_navigation_results(generated_dir: str, reference_dir: str) -> Tuple[bool, str]:
+    """Original JSON trajectory sequence comparison (unchanged logic).
+
+    NOTE: Do not modify this function's core behavior. Image comparison is handled by a separate
+    wrapper function `compare_navigation_and_images` to avoid side effects on existing tests.
+    """
+    # --- Enhanced logic ---
+    # To support both "caller passes seq_path root directory" and "legacy call (leaf trajectory directory)" forms,
+    # here we use a symmetric data.json discovery strategy for both generated and reference sides:
+    # 1. If the current directory directly contains data.json, use that file.
+    # 2. Otherwise, traverse one level down into subdirectories (sorted alphabetically), looking for <dir>/data.json.
+    # 3. Otherwise, search within two nested levels (dir/subdir/data.json) and use the first match found.
+    # 4. If not found, report an error. This is compatible with the legacy "generated=root, reference=leaf" usage,
+    #    and also allows both sides to provide the root directory.
+
+    if not os.path.isdir(generated_dir):
+        return False, f"Generated directory does not exist or is not a directory: {generated_dir}"
+    if not os.path.isdir(reference_dir):
+        return False, f"Reference directory does not exist or is not a directory: {reference_dir}"
+
+    try:
+        generated_file = _locate_first_data_json(generated_dir)
+        if generated_file is None:
+            return False, f"Could not locate data.json under generated directory: {generated_dir}"
+    except Exception as e:  # pylint: disable=broad-except
+        return False, f"Error locating generated data file: {e}"
+
+    try:
+        reference_file = _locate_first_data_json(reference_dir)
+        if reference_file is None:
+            # To preserve legacy behavior, if reference_dir/data.json exists but was not detected above (should not happen in theory), check once more
+            candidate = os.path.join(reference_dir, "data.json")
+            if os.path.isfile(candidate):
+                reference_file = candidate
+            else:
+                return False, f"Could not locate data.json under reference directory: {reference_dir}"
+    except Exception as e:  # pylint: disable=broad-except
+        return False, f"Error locating reference data file: {e}"
+
+    return compare_trajectory_sequences(generated_file, reference_file)
+
+
+def compare_navigation_and_images(
+    generated_seq_dir: str,
+    reference_seq_dir: str,
+    generated_root_for_images: Optional[str] = None,
+    reference_root_for_images: Optional[str] = None,
+    rgb_abs_tolerance: int = 0,
+    depth_abs_tolerance: float = 0.0,
+    allowed_rgb_diff_ratio: float = 0.0,
+    allowed_depth_diff_ratio: float = 0.5,
+    depth_scale_auto: bool = False,
+    fail_if_images_missing: bool = False,
+) -> Tuple[bool, str]:
+    """Wrapper that preserves original JSON comparison while optionally adding first-frame image comparison.
+
+    Args:
+        generated_seq_dir: Path to generated seq_path root used by original comparator.
+        reference_seq_dir: Path to reference seq_path root.
+        generated_root_for_images: Root (parent of obs_path) or the obs_path itself for generated images.
+        reference_root_for_images: Same as above for reference. If None, image comparison may be skipped.
+        rgb_abs_tolerance: RGB absolute per-channel tolerance.
+        depth_abs_tolerance: Depth absolute tolerance.
+        allowed_rgb_diff_ratio: Allowed differing RGB pixel ratio.
+        allowed_depth_diff_ratio: Allowed differing depth pixel ratio.
+        depth_scale_auto: Auto scale depth if uint16 millimeters.
+        fail_if_images_missing: If True, treat missing obs_path as failure; otherwise skip.
+
+    Returns:
+        (success, message) combined result.
+    """
+    traj_ok, traj_msg = compare_navigation_results(generated_seq_dir, reference_seq_dir)
+
+    # Determine image roots; default to parent of seq_dir if not explicitly provided
+    gen_img_root = generated_root_for_images or os.path.dirname(generated_seq_dir.rstrip(os.sep))
+    ref_img_root = reference_root_for_images or os.path.dirname(reference_seq_dir.rstrip(os.sep))
+
+    img_ok = True
+    img_msg = "image comparison skipped"
+
+    if generated_root_for_images is not None or reference_root_for_images is not None:
+        # User explicitly passed at least one root -> attempt compare
+        img_ok, img_msg = compare_first_frame_images(
+            generated_root=gen_img_root,
+            reference_root=ref_img_root,
+            rgb_abs_tolerance=rgb_abs_tolerance,
+            depth_abs_tolerance=depth_abs_tolerance,
+            allowed_rgb_diff_ratio=allowed_rgb_diff_ratio,
+            allowed_depth_diff_ratio=allowed_depth_diff_ratio,
+            depth_scale_auto=depth_scale_auto,
+        )
+    else:
+        # Implicit attempt only if both obs_path exist under parent paths
+        gen_obs_candidate = os.path.join(gen_img_root, "obs_path")
+        ref_obs_candidate = os.path.join(ref_img_root, "obs_path")
+        if os.path.isdir(gen_obs_candidate) and os.path.isdir(ref_obs_candidate):
+            img_ok, img_msg = compare_first_frame_images(
+                generated_root=gen_img_root,
+                reference_root=ref_img_root,
+                rgb_abs_tolerance=rgb_abs_tolerance,
+                depth_abs_tolerance=depth_abs_tolerance,
+                allowed_rgb_diff_ratio=allowed_rgb_diff_ratio,
+                allowed_depth_diff_ratio=allowed_depth_diff_ratio,
+                depth_scale_auto=depth_scale_auto,
+            )
+        else:
+            if fail_if_images_missing:
+                missing = []
+                if not os.path.isdir(gen_obs_candidate):
+                    missing.append(f"generated:{gen_obs_candidate}")
+                if not os.path.isdir(ref_obs_candidate):
+                    missing.append(f"reference:{ref_obs_candidate}")
+                img_ok = False
+                img_msg = "obs_path missing -> " + ", ".join(missing)
+            else:
+                img_msg = "obs_path not found in one or both roots; skipped"
+
+    overall = traj_ok and img_ok
+    message = f"trajectory: {traj_msg}; images: {img_msg}"
+    return overall, message if overall else f"Mismatch - {message}"
+
+
+def compare_trajectory_sequences(generated_file: str, reference_file: str, tolerance: float = 1e-6) -> Tuple[bool, str]:
+    """
+    Compare trajectory sequence files with numerical tolerance.
+
+    Args:
+        generated_file: Path to generated trajectory file
+        reference_file: Path to reference trajectory file
+        tolerance: Numerical tolerance for floating point comparisons
+
+    Returns:
+        Tuple[bool, str]: (success, message)
+    """
+    try:
+        # Check if files exist
+        if not os.path.exists(generated_file):
+            return False, f"Generated file does not exist: {generated_file}"
+
+        if not os.path.exists(reference_file):
+            return False, f"Reference file does not exist: {reference_file}"
+
+        # Load JSON files
+        with open(generated_file, 'r') as f:
+            generated_data = json.load(f)
+
+        with open(reference_file, 'r') as f:
+            reference_data = json.load(f)
+
+        # Compare the JSON structures
+        success, message = _compare_data_structures(generated_data, reference_data, tolerance)
+
+        if success:
+            return True, "Trajectory sequences match within tolerance"
+        else:
+            return False, f"Trajectory sequences differ: {message}"
+
+    except json.JSONDecodeError as e:
+        return False, f"JSON decode error: {e}"
+    except Exception as e:
+        return False, f"Error comparing trajectory sequences: {e}"
+
+
+def _compare_data_structures(data1: Any, data2: Any, tolerance: float, path: str = "") -> Tuple[bool, str]:
+    """
+    Recursively compare two data structures with numerical tolerance.
+
+    Args:
+        data1: First data structure
+        data2: Second data structure
+        tolerance: Numerical tolerance for floating point comparisons
+        path: Current path in the data structure for error reporting
+
+    Returns:
+        Tuple[bool, str]: (success, error_message)
+    """
+    # Check if types are the same
+    if type(data1) != type(data2):
+        return False, f"Type mismatch at {path}: {type(data1)} vs {type(data2)}"
+
+    # Handle dictionaries
+    if isinstance(data1, dict):
+        if set(data1.keys()) != set(data2.keys()):
+            return False, f"Key mismatch at {path}: {set(data1.keys())} vs {set(data2.keys())}"
+
+        for key in data1.keys():
+            new_path = f"{path}.{key}" if path else key
+            success, message = _compare_data_structures(data1[key], data2[key], tolerance, new_path)
+            if not success:
+                return False, message
+
+    # Handle lists
+    elif isinstance(data1, list):
+        if len(data1) != len(data2):
+            return False, f"List length mismatch at {path}: {len(data1)} vs {len(data2)}"
+
+        for i, (item1, item2) in enumerate(zip(data1, data2)):
+            new_path = f"{path}[{i}]" if path else f"[{i}]"
+            success, message = _compare_data_structures(item1, item2, tolerance, new_path)
+            if not success:
+                return False, message
+
+    # Handle numerical values
+    elif isinstance(data1, (int, float)):
+        if isinstance(data2, (int, float)):
+            if abs(data1 - data2) > tolerance:
+                return (
+                    False,
+                    f"Numerical difference at {path}: {data1} vs {data2} (diff: {abs(data1 - data2)}, tolerance: {tolerance})",
+                )
+        else:
+            return False, f"Type mismatch at {path}: number vs {type(data2)}"
+
+    # Handle strings and other exact comparison types
+    elif isinstance(data1, (str, bool, type(None))):
+        if data1 != data2:
+            return False, f"Value mismatch at {path}: {data1} vs {data2}"
+
+    # Handle unknown types
+    else:
+        if data1 != data2:
+            return False, f"Value mismatch at {path}: {data1} vs {data2}"
+
+    return True, ""
+
+def _locate_first_data_json(root: str) -> Optional[str]:
+    """Locate a data.json file under root with a shallow, deterministic strategy.
+
+    Strategy (stop at first match to keep behavior predictable & lightweight):
+      1. If root/data.json exists -> return it.
+      2. Enumerate immediate subdirectories (sorted). For each d:
+         - if d/data.json exists -> return it.
+      3. Enumerate immediate subdirectories again; for each d enumerate its subdirectories (sorted) and
+         look for d/sub/data.json -> return first match.
+      4. If none found -> return None.
+    """
+    # 1. root/data.json
+    candidate = os.path.join(root, "data.json")
+    if os.path.isfile(candidate):
+        return candidate
+
+    try:
+        first_level = [d for d in os.listdir(root) if os.path.isdir(os.path.join(root, d))]
+    except FileNotFoundError:
+        return None
+
+    first_level.sort()
+
+    # 2. d/data.json
+    for d in first_level:
+        c = os.path.join(root, d, "data.json")
+        if os.path.isfile(c):
+            return c
+
+    # 3. d/sub/data.json
+    for d in first_level:
+        d_path = os.path.join(root, d)
+        try:
+            second_level = [s for s in os.listdir(d_path) if os.path.isdir(os.path.join(d_path, s))]
+        except FileNotFoundError:
+            continue
+        second_level.sort()
+        for s in second_level:
+            c = os.path.join(d_path, s, "data.json")
+            if os.path.isfile(c):
+                return c
+
+    return None
+
+
+def compare_first_frame_images(
+    generated_root: str,
+    reference_root: str,
+    rgb_dir_name: str = "rgb",
+    depth_dir_name: str = "depth",
+    scene_dir: Optional[str] = None,
+    traj_dir: Optional[str] = None,
+    rgb_abs_tolerance: int = 0,
+    depth_abs_tolerance: float = 0.0,
+    allowed_rgb_diff_ratio: float = 0.0,
+    allowed_depth_diff_ratio: float = 0.0,
+    compute_psnr: bool = True,
+    compute_mse: bool = True,
+    depth_scale_auto: bool = False,
+) -> Tuple[bool, str]:
+    """Compare only the first frame (index 0) of RGB & depth images between generated and reference.
+
+    This is a lightweight check to validate pipeline correctness without scanning all frames.
+
+    Args:
+        generated_root: Path to generated run root (may contain `obs_path` or be the `obs_path`).
+        reference_root: Path to reference run root (same structure as generated_root).
+        rgb_dir_name: Subdirectory name for RGB frames under a trajectory directory.
+        depth_dir_name: Subdirectory name for depth frames under a trajectory directory.
+        scene_dir: Optional explicit scene directory name (e.g. "6f"); if None will auto-pick first.
+        traj_dir: Optional explicit trajectory directory (e.g. "0"); if None will auto-pick first.
+        rgb_abs_tolerance: Per-channel absolute pixel tolerance (0 requires exact match).
+        depth_abs_tolerance: Absolute tolerance for depth value differences (after optional scaling).
+        allowed_rgb_diff_ratio: Max allowed ratio of differing RGB pixels (0.01 -> 1%).
+        allowed_depth_diff_ratio: Max allowed ratio of differing depth pixels beyond tolerance.
+        compute_psnr: Whether to compute PSNR metric for reporting.
+        compute_mse: Whether to compute MSE metric for reporting.
+        depth_scale_auto: If True, attempt simple heuristic scaling for uint16 depth (divide by 1000 if max > 10000).
+
+    Returns:
+        (success, message) summary of comparison.
+    """
+    try:
+        gen_obs = _resolve_obs_path(generated_root)
+        ref_obs = _resolve_obs_path(reference_root)
+        if gen_obs is None:
+            return False, f"Cannot locate obs_path under generated root: {generated_root}"
+        if ref_obs is None:
+            return False, f"Cannot locate obs_path under reference root: {reference_root}"
+
+        scene_dir = scene_dir or _pick_first_subdir(gen_obs)
+        if scene_dir is None:
+            return False, f"No scene directory found in {gen_obs}"
+        ref_scene_dir = scene_dir if os.path.isdir(os.path.join(ref_obs, scene_dir)) else _pick_first_subdir(ref_obs)
+        if ref_scene_dir is None:
+            return False, f"No matching scene directory in reference: {ref_obs}"
+
+        gen_scene_path = os.path.join(gen_obs, scene_dir)
+        ref_scene_path = os.path.join(ref_obs, ref_scene_dir)
+
+        traj_dir = traj_dir or _pick_first_subdir(gen_scene_path)
+        if traj_dir is None:
+            return False, f"No trajectory directory in {gen_scene_path}"
+        ref_traj_dir = (
+            traj_dir if os.path.isdir(os.path.join(ref_scene_path, traj_dir)) else _pick_first_subdir(ref_scene_path)
+        )
+        if ref_traj_dir is None:
+            return False, f"No trajectory directory in reference scene path {ref_scene_path}"
+
+        gen_traj_path = os.path.join(gen_scene_path, traj_dir)
+        ref_traj_path = os.path.join(ref_scene_path, ref_traj_dir)
+
+        # RGB comparison
+        rgb_result, rgb_msg = _compare_single_frame_rgb(
+            gen_traj_path,
+            ref_traj_path,
+            rgb_dir_name,
+            rgb_abs_tolerance,
+            allowed_rgb_diff_ratio,
+            compute_psnr,
+            compute_mse,
+        )
+
+        # Depth comparison (optional if depth folder exists)
+        depth_result, depth_msg = _compare_single_frame_depth(
+            gen_traj_path,
+            ref_traj_path,
+            depth_dir_name,
+            depth_abs_tolerance,
+            allowed_depth_diff_ratio,
+            compute_psnr,
+            compute_mse,
+            depth_scale_auto,
+        )
+
+        success = rgb_result and depth_result
+        combined_msg = f"RGB: {rgb_msg}; Depth: {depth_msg}"
+        return success, ("Images match - " + combined_msg) if success else ("Image mismatch - " + combined_msg)
+    except Exception as e:  # pylint: disable=broad-except
+        return False, f"Error during first-frame image comparison: {e}"
+
+
+def _resolve_obs_path(root: str) -> Optional[str]:
+    """Return the obs_path directory. Accept either the root itself or its child."""
+    if not os.path.isdir(root):
+        return None
+    if os.path.basename(root) == "obs_path":
+        return root
+    candidate = os.path.join(root, "obs_path")
+    return candidate if os.path.isdir(candidate) else None
+
+
+def _pick_first_subdir(parent: str) -> Optional[str]:
+    """Pick the first alphanumerically sorted subdirectory name under parent."""
+    try:
+        subs = [d for d in os.listdir(parent) if os.path.isdir(os.path.join(parent, d))]
+        if not subs:
+            return None
+        subs.sort()
+        return subs[0]
+    except FileNotFoundError:
+        return None
+
+
+def _find_first_frame_file(folder: str, exts: Tuple[str, ...]) -> Optional[str]:
+    """Find the smallest numeral file with one of extensions; returns absolute path."""
+    if not os.path.isdir(folder):
+        return None
+    candidates = []
+    for f in os.listdir(folder):
+        lower = f.lower()
+        for e in exts:
+            if lower.endswith(e):
+                num_part = os.path.splitext(f)[0]
+                if num_part.isdigit():
+                    candidates.append((int(num_part), f))
+                elif f.startswith("0"):  # fallback for names like 0.jpg
+                    candidates.append((0, f))
+                break
+    if not candidates:
+        return None
+    candidates.sort(key=lambda x: x[0])
+    return os.path.join(folder, candidates[0][1])
+
+
+def _compare_single_frame_rgb(
+    gen_traj_path: str,
+    ref_traj_path: str,
+    rgb_dir_name: str,
+    abs_tol: int,
+    allowed_ratio: float,
+    compute_psnr: bool,
+    compute_mse: bool,
+) -> Tuple[bool, str]:
+    rgb_gen_dir = os.path.join(gen_traj_path, rgb_dir_name)
+    rgb_ref_dir = os.path.join(ref_traj_path, rgb_dir_name)
+    if not os.path.isdir(rgb_gen_dir) or not os.path.isdir(rgb_ref_dir):
+        return (
+            False,
+            f"RGB directory missing (generated: {os.path.isdir(rgb_gen_dir)}, reference: {os.path.isdir(rgb_ref_dir)})",
+        )
+
+    gen_file = _find_first_frame_file(rgb_gen_dir, (".jpg", ".png", ".jpeg"))
+    ref_file = _find_first_frame_file(rgb_ref_dir, (".jpg", ".png", ".jpeg"))
+    if not gen_file or not ref_file:
+        return False, "First RGB frame file not found in one of the directories"
+
+    gen_img = cv2.imread(gen_file, cv2.IMREAD_COLOR)
+    ref_img = cv2.imread(ref_file, cv2.IMREAD_COLOR)
+    if gen_img is None or ref_img is None:
+        return False, "Failed to read RGB images"
+    if gen_img.shape != ref_img.shape:
+        return False, f"RGB shape mismatch {gen_img.shape} vs {ref_img.shape}"
+
+    diff = np.abs(gen_img.astype(np.int16) - ref_img.astype(np.int16))
+    diff_mask = np.any(diff > abs_tol, axis=2)
+    diff_ratio = float(diff_mask.sum()) / diff_mask.size
+
+    metrics_parts = [f"diff_pixels_ratio={diff_ratio:.4f}"]
+    flag = False
+    if compute_mse or compute_psnr:
+        mse = float((diff**2).mean())
+        if compute_mse:
+            metrics_parts.append(f"mse={mse:.2f}")
+        if compute_psnr:
+            if mse == 0.0:
+                psnr = float('inf')
+                flag = True
+            else:
+                psnr = 10.0 * math.log10((255.0**2) / mse)
+            if math.isinf(psnr):
+                metrics_parts.append("psnr=inf")
+                flag = True
+            else:
+                metrics_parts.append(f"psnr={psnr:.2f}dB")
+                if psnr >= 40.0:
+                    flag = True
+
+    passed = diff_ratio <= allowed_ratio or flag
+    status = "OK" if passed else "FAIL"
+    return passed, f"{status} (abs_tol={abs_tol}, allowed_ratio={allowed_ratio}, {' '.join(metrics_parts)})"
+
+
+def _compare_single_frame_depth(
+    gen_traj_path: str,
+    ref_traj_path: str,
+    depth_dir_name: str,
+    abs_tol: float,
+    allowed_ratio: float,
+    compute_psnr: bool,
+    compute_mse: bool,
+    auto_scale: bool,
+) -> Tuple[bool, str]:
+    depth_gen_dir = os.path.join(gen_traj_path, depth_dir_name)
+    depth_ref_dir = os.path.join(ref_traj_path, depth_dir_name)
+    if not os.path.isdir(depth_gen_dir) or not os.path.isdir(depth_ref_dir):
+        return (
+            False,
+            f"Depth directory missing (generated: {os.path.isdir(depth_gen_dir)}, reference: {os.path.isdir(depth_ref_dir)})",
+        )
+
+    gen_file = _find_first_frame_file(depth_gen_dir, (".png", ".exr"))
+    ref_file = _find_first_frame_file(depth_ref_dir, (".png", ".exr"))
+    if not gen_file or not ref_file:
+        return False, "First depth frame file not found in one of the directories"
+
+    gen_img = cv2.imread(gen_file, cv2.IMREAD_UNCHANGED)
+    ref_img = cv2.imread(ref_file, cv2.IMREAD_UNCHANGED)
+    if gen_img is None or ref_img is None:
+        return False, "Failed to read depth images"
+    if gen_img.shape != ref_img.shape:
+        return False, f"Depth shape mismatch {gen_img.shape} vs {ref_img.shape}"
+
+    gen_depth = _prepare_depth_array(gen_img, auto_scale)
+    ref_depth = _prepare_depth_array(ref_img, auto_scale)
+    if gen_depth.shape != ref_depth.shape:
+        return False, f"Depth array shape mismatch {gen_depth.shape} vs {ref_depth.shape}"
+
+    diff = np.abs(gen_depth - ref_depth)
+    diff_mask = diff > abs_tol
+    diff_ratio = float(diff_mask.sum()) / diff_mask.size
+
+    metrics_parts = [f"diff_pixels_ratio={diff_ratio:.4f}"]
+    if compute_mse or compute_psnr:
+        mse = float((diff**2).mean())
+        if compute_mse:
+            metrics_parts.append(f"mse={mse:.6f}")
+        if compute_psnr:
+            # Estimate dynamic range from reference depth
+            dr = float(ref_depth.max() - ref_depth.min()) or 1.0
+            if mse == 0.0:
+                psnr = float('inf')
+            else:
+                psnr = 10.0 * math.log10((dr**2) / mse)
+            metrics_parts.append("psnr=inf" if math.isinf(psnr) else f"psnr={psnr:.2f}dB")
+
+    passed = diff_ratio <= allowed_ratio
+    status = "OK" if passed else "FAIL"
+    return passed, f"{status} (abs_tol={abs_tol}, allowed_ratio={allowed_ratio}, {' '.join(metrics_parts)})"
+
+
+def _prepare_depth_array(arr: np.ndarray, auto_scale: bool) -> np.ndarray:
+    """Convert raw depth image to float32 array; apply simple heuristic scaling if requested."""
+    if arr.dtype == np.uint16:
+        depth = arr.astype(np.float32)
+        if auto_scale and depth.max() > 10000:  # likely millimeters
+            depth /= 1000.0
+        return depth
+    if arr.dtype == np.float32:
+        return arr
+    # Fallback: convert to float
+    return arr.astype(np.float32)

tests/integration/data_comparators/simbox_comparator.py

@@ -0,0 +1,398 @@
+"""
+Simbox Output Comparator
+
+This module provides functionality to compare two Simbox task output directories.
+It compares both meta_info.pkl and LMDB database contents, handling different data types:
+- JSON data (dict/list)
+- Scalar data (numerical arrays/lists)
+- Image data (encoded images)
+- Proprioception data (joint states, gripper states)
+- Object data (object poses and properties)
+- Action data (robot actions)
+"""
+
+import argparse
+import pickle
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+import cv2
+import lmdb
+import numpy as np
+
+
+class SimboxComparator:
+    """Comparator for Simbox task output directories."""
+
+    def __init__(self, dir1: str, dir2: str, tolerance: float = 1e-6, image_psnr_threshold: float = 30.0):
+        """
+        Initialize the comparator.
+
+        Args:
+            dir1: Path to the first output directory
+            dir2: Path to the second output directory
+            tolerance: Numerical tolerance for floating point comparisons
+            image_psnr_threshold: PSNR threshold (dB) for considering images as acceptable match
+        """
+        self.dir1 = Path(dir1)
+        self.dir2 = Path(dir2)
+        self.tolerance = tolerance
+        self.image_psnr_threshold = image_psnr_threshold
+        self.mismatches = []
+        self.warnings = []
+        self.image_psnr_values: List[float] = []
+
+    def load_directory(self, directory: Path) -> Tuple[Optional[Dict], Optional[Any], Optional[Any]]:
+        """Load meta_info.pkl and LMDB database from directory."""
+        meta_path = directory / "meta_info.pkl"
+        lmdb_path = directory / "lmdb"
+
+        if not directory.is_dir() or not meta_path.exists() or not lmdb_path.is_dir():
+            print(f"Error: '{directory}' is not a valid output directory.")
+            print("It must contain 'meta_info.pkl' and an 'lmdb' subdirectory.")
+            return None, None, None
+
+        try:
+            with open(meta_path, "rb") as f:
+                meta_info = pickle.load(f)
+        except Exception as e:
+            print(f"Error loading meta_info.pkl from {directory}: {e}")
+            return None, None, None
+
+        try:
+            env = lmdb.open(str(lmdb_path), readonly=True, lock=False, readahead=False, meminit=False)
+            txn = env.begin(write=False)
+        except Exception as e:
+            print(f"Error opening LMDB database at {lmdb_path}: {e}")
+            return None, None, None
+
+        return meta_info, txn, env
+
+    def compare_metadata(self, meta1: Dict, meta2: Dict) -> bool:
+        """Compare high-level metadata."""
+        identical = True
+
+        if meta1.get("num_steps") != meta2.get("num_steps"):
+            self.mismatches.append(f"num_steps differ: {meta1.get('num_steps')} vs {meta2.get('num_steps')}")
+            identical = False
+
+        return identical
+
+    def get_key_categories(self, meta: Dict) -> Dict[str, set]:
+        """Extract key categories from metadata."""
+        key_to_category = {}
+        for category, keys in meta.get("keys", {}).items():
+            for key in keys:
+                key_bytes = key if isinstance(key, bytes) else key.encode()
+                key_to_category[key_bytes] = category
+
+        return key_to_category
+
+    def compare_json_data(self, key: bytes, data1: Any, data2: Any) -> bool:
+        """Compare JSON/dict/list data."""
+        if type(data1) != type(data2):
+            self.mismatches.append(f"[{key.decode()}] Type mismatch: {type(data1).__name__} vs {type(data2).__name__}")
+            return False
+
+        if isinstance(data1, dict):
+            if set(data1.keys()) != set(data2.keys()):
+                self.mismatches.append(f"[{key.decode()}] Dict keys differ")
+                return False
+            for k in data1.keys():
+                if not self.compare_json_data(key, data1[k], data2[k]):
+                    return False
+        elif isinstance(data1, list):
+            if len(data1) != len(data2):
+                self.mismatches.append(f"[{key.decode()}] List length differ: {len(data1)} vs {len(data2)}")
+                return False
+            # For lists, compare sample elements to avoid excessive output
+            if len(data1) > 10:
+                sample_indices = [0, len(data1) // 2, -1]
+                for idx in sample_indices:
+                    if not self.compare_json_data(key, data1[idx], data2[idx]):
+                        return False
+            else:
+                for i, (v1, v2) in enumerate(zip(data1, data2)):
+                    if not self.compare_json_data(key, v1, v2):
+                        return False
+        else:
+            if data1 != data2:
+                self.mismatches.append(f"[{key.decode()}] Value mismatch: {data1} vs {data2}")
+                return False
+
+        return True
+
+    def compare_numerical_data(self, key: bytes, data1: Any, data2: Any) -> bool:
+        """Compare numerical data (arrays, lists of numbers)."""
+        # Convert to numpy arrays for comparison
+        try:
+            if isinstance(data1, list):
+                arr1 = np.array(data1)
+                arr2 = np.array(data2)
+            else:
+                arr1 = data1
+                arr2 = data2
+
+            if arr1.shape != arr2.shape:
+                self.mismatches.append(f"[{key.decode()}] Shape mismatch: {arr1.shape} vs {arr2.shape}")
+                return False
+
+            if not np.allclose(arr1, arr2, rtol=self.tolerance, atol=self.tolerance):
+                diff = np.abs(arr1 - arr2)
+                max_diff = np.max(diff)
+                mean_diff = np.mean(diff)
+                self.mismatches.append(
+                    f"[{key.decode()}] Numerical difference: max={max_diff:.6e}, mean={mean_diff:.6e}"
+                )
+                return False
+
+        except Exception as e:
+            self.warnings.append(f"[{key.decode()}] Error comparing numerical data: {e}")
+            return False
+
+        return True
+
+    def compare_image_data(self, key: bytes, data1: np.ndarray, data2: np.ndarray) -> bool:
+        """Compare image data (encoded as uint8 arrays)."""
+        try:
+            # Decode images
+            img1 = cv2.imdecode(data1, cv2.IMREAD_UNCHANGED)
+            img2 = cv2.imdecode(data2, cv2.IMREAD_UNCHANGED)
+
+            if img1 is None or img2 is None:
+                self.warnings.append(f"[{key.decode()}] Could not decode image, using binary comparison")
+                return np.array_equal(data1, data2)
+
+            # Compare shapes
+            if img1.shape != img2.shape:
+                self.mismatches.append(f"[{key.decode()}] Image shape mismatch: {img1.shape} vs {img2.shape}")
+                return False
+
+            # Calculate PSNR for tracking average quality
+            if np.array_equal(img1, img2):
+                psnr = 100.0
+            else:
+                diff_float = img1.astype(np.float32) - img2.astype(np.float32)
+                mse = np.mean(diff_float ** 2)
+                if mse == 0:
+                    psnr = 100.0
+                else:
+                    max_pixel = 255.0
+                    psnr = 20 * np.log10(max_pixel / np.sqrt(mse))
+
+            self.image_psnr_values.append(psnr)
+
+            try:
+                print(f"[{key.decode()}] PSNR: {psnr:.2f} dB")
+            except Exception:
+                print(f"[<binary key>] PSNR: {psnr:.2f} dB")
+
+            return True
+
+        except Exception as e:
+            self.warnings.append(f"[{key.decode()}] Error comparing image: {e}")
+            return False
+
+    def _save_comparison_image(self, key: bytes, img1: np.ndarray, img2: np.ndarray, diff: np.ndarray):
+        """Save comparison visualization for differing images."""
+        try:
+            output_dir = Path("image_comparisons")
+            output_dir.mkdir(exist_ok=True)
+
+            # Normalize difference for visualization
+            if len(diff.shape) == 3:
+                diff_vis = np.clip(diff * 5, 0, 255).astype(np.uint8)
+            else:
+                diff_vis = np.clip(diff * 5, 0, 255).astype(np.uint8)
+
+            # Ensure RGB format for concatenation
+            def ensure_rgb(img):
+                if len(img.shape) == 2:
+                    return cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
+                elif len(img.shape) == 3 and img.shape[2] == 4:
+                    return cv2.cvtColor(img, cv2.COLOR_BGRA2RGB)
+                return img
+
+            img1_rgb = ensure_rgb(img1)
+            img2_rgb = ensure_rgb(img2)
+            diff_rgb = ensure_rgb(diff_vis)
+
+            # Concatenate horizontally
+            combined = np.hstack([img1_rgb, img2_rgb, diff_rgb])
+
+            # Save
+            safe_key = key.decode().replace("/", "_").replace("\\", "_").replace(":", "_")
+            output_path = output_dir / f"diff_{safe_key}.png"
+            cv2.imwrite(str(output_path), cv2.cvtColor(combined, cv2.COLOR_RGB2BGR))
+
+        except Exception as e:
+            self.warnings.append(f"Failed to save comparison image for {key.decode()}: {e}")
+
+    def compare_value(self, key: bytes, category: str, val1: bytes, val2: bytes) -> bool:
+        """Compare a single key-value pair based on its category."""
+        # Output the category for the current key being compared
+        try:
+            print(f"[{key.decode()}] category: {category}")
+        except Exception:
+            print(f"[<binary key>] category: {category}")
+
+        if val1 is None and val2 is None:
+            return True
+        if val1 is None or val2 is None:
+            self.mismatches.append(f"[{key.decode()}] Key exists in one dataset but not the other")
+            return False
+
+        try:
+            data1 = pickle.loads(val1)
+            data2 = pickle.loads(val2)
+        except Exception as e:
+            self.warnings.append(f"[{key.decode()}] Error unpickling data: {e}")
+            return val1 == val2
+
+        # Route to appropriate comparison based on category
+        if category == "json_data":
+            return self.compare_json_data(key, data1, data2)
+        elif category in ["scalar_data", "proprio_data", "object_data", "action_data"]:
+            return self.compare_numerical_data(key, data1, data2)
+        elif category.startswith("images."):
+            # Image data is stored as numpy uint8 array
+            if isinstance(data1, np.ndarray) and isinstance(data2, np.ndarray):
+                return self.compare_image_data(key, data1, data2)
+            else:
+                self.warnings.append(f"[{key.decode()}] Expected numpy array for image data")
+                return False
+        else:
+            # Unknown category, try generic comparison
+            self.warnings.append(f"[{key.decode()}] Unknown category '{category}', using binary comparison")
+            return val1 == val2
+
+    def compare(self) -> bool:
+        """Execute full comparison."""
+        print(f"Comparing directories:")
+        print(f"  Dir1: {self.dir1}")
+        print(f"  Dir2: {self.dir2}\n")
+
+        # Load data
+        meta1, txn1, env1 = self.load_directory(self.dir1)
+        meta2, txn2, env2 = self.load_directory(self.dir2)
+
+        if meta1 is None or meta2 is None:
+            print("Aborting due to loading errors.")
+            return False
+
+        print("Successfully loaded data from both directories.\n")
+
+        # Compare metadata
+        print("Comparing metadata...")
+        self.compare_metadata(meta1, meta2)
+
+        # Get key categories
+        key_cat1 = self.get_key_categories(meta1)
+        key_cat2 = self.get_key_categories(meta2)
+
+        keys1 = set(key_cat1.keys())
+        keys2 = set(key_cat2.keys())
+
+        # Check key sets
+        if keys1 != keys2:
+            missing_in_2 = sorted([k.decode() for k in keys1 - keys2])
+            missing_in_1 = sorted([k.decode() for k in keys2 - keys1])
+            if missing_in_2:
+                self.mismatches.append(f"Keys missing in dir2: {missing_in_2[:10]}")
+            if missing_in_1:
+                self.mismatches.append(f"Keys missing in dir1: {missing_in_1[:10]}")
+
+        # Compare common keys
+        common_keys = sorted(list(keys1.intersection(keys2)))
+        print(f"Comparing {len(common_keys)} common keys...\n")
+
+        for i, key in enumerate(common_keys):
+            if i % 100 == 0 and i > 0:
+                print(f"Progress: {i}/{len(common_keys)} keys compared...")
+
+            category = key_cat1.get(key, "unknown")
+            val1 = txn1.get(key)
+            val2 = txn2.get(key)
+
+            self.compare_value(key, category, val1, val2)
+
+        if self.image_psnr_values:
+            avg_psnr = sum(self.image_psnr_values) / len(self.image_psnr_values)
+            print(
+                f"\nImage PSNR average over {len(self.image_psnr_values)} images: "
+                f"{avg_psnr:.2f} dB (threshold {self.image_psnr_threshold:.2f} dB)"
+            )
+            if avg_psnr < self.image_psnr_threshold:
+                self.mismatches.append(
+                    f"Average image PSNR {avg_psnr:.2f} dB below threshold "
+                    f"{self.image_psnr_threshold:.2f} dB"
+                )
+        else:
+            print("\nNo image entries found for PSNR calculation.")
+
+        # Cleanup
+        env1.close()
+        env2.close()
+
+        # Print results
+        print("\n" + "=" * 80)
+        print("COMPARISON RESULTS")
+        print("=" * 80)
+
+        if self.warnings:
+            print(f"\nWarnings ({len(self.warnings)}):")
+            for warning in self.warnings[:20]:
+                print(f"  - {warning}")
+            if len(self.warnings) > 20:
+                print(f"  ... and {len(self.warnings) - 20} more warnings")
+
+        if self.mismatches:
+            print(f"\nMismatches found ({len(self.mismatches)}):")
+            for mismatch in self.mismatches[:30]:
+                print(f"  - {mismatch}")
+            if len(self.mismatches) > 30:
+                print(f"  ... and {len(self.mismatches) - 30} more mismatches")
+            print("\n❌ RESULT: Directories are DIFFERENT")
+            return False
+        else:
+            print("\n✅ RESULT: Directories are IDENTICAL")
+            return True
+
+
+def main():
+    """Main entry point."""
+    parser = argparse.ArgumentParser(
+        description="Compare two Simbox task output directories.",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  %(prog)s --dir1 output/run1 --dir2 output/run2
+  %(prog)s --dir1 output/run1 --dir2 output/run2 --tolerance 1e-5
+  %(prog)s --dir1 output/run1 --dir2 output/run2 --image-psnr 40.0
+        """,
+    )
+    parser.add_argument("--dir1", type=str, required=True, help="Path to the first output directory")
+    parser.add_argument("--dir2", type=str, required=True, help="Path to the second output directory")
+    parser.add_argument(
+        "--tolerance",
+        type=float,
+        default=1e-6,
+        help="Numerical tolerance for floating point comparisons (default: 1e-6)",
+    )
+    parser.add_argument(
+        "--image-psnr",
+        type=float,
+        default=37.0,
+        help="PSNR threshold (dB) for considering images as matching (default: 37.0)",
+    )
+
+    args = parser.parse_args()
+
+    comparator = SimboxComparator(args.dir1, args.dir2, tolerance=args.tolerance, image_psnr_threshold=args.image_psnr)
+
+    success = comparator.compare()
+    exit(0 if success else 1)
+
+
+if __name__ == "__main__":
+    main()

tests/integration/simbox/runners/simbox_pipeline_runner.py

@@ -0,0 +1,30 @@
+"""
+simbox pipeline runner for Isaac Sim subprocess execution.
+This script runs in the Isaac Sim environment and is called by the main test.
+"""
+
+
+import sys
+
+from tests.integration.base.test_harness import IntegrationTestHarness
+
+# Add paths to sys.path
+sys.path.append("./")
+sys.path.append("./data_engine")
+sys.path.append("./tests/integration")
+
+
+def run_simbox_pipeline():
+    """
+    Run the simbox pipeline test in Isaac Sim environment.
+    """
+    harness = IntegrationTestHarness(config_path="configs/simbox/de_pipe_template.yaml", seed=42, random_num=1)
+
+    # Run the pipeline and validate output
+    harness.run_test_end_to_end(min_output_files=6)
+
+    print("✓ simbox pipeline test completed successfully")
+
+
+if __name__ == "__main__":
+    run_simbox_pipeline()

tests/integration/simbox/runners/simbox_plan_runner.py

@@ -0,0 +1,29 @@
+"""
+simbox plan runner for Isaac Sim subprocess execution.
+This script runs in the Isaac Sim environment and is called by the main test.
+"""
+
+import sys
+
+from tests.integration.base.test_harness import IntegrationTestHarness
+
+# Add paths to sys.path
+sys.path.append("./")
+sys.path.append("./data_engine")
+sys.path.append("./tests/integration")
+
+
+def run_simbox_plan():
+    """
+    Run the simbox plan test in Isaac Sim environment.
+    """
+    harness = IntegrationTestHarness(config_path="configs/simbox/de_plan_template.yaml", seed=42, random_num=1)
+
+    # Run the pipeline and validate output
+    harness.run_test_end_to_end(min_output_files=1)
+
+    print("✓ simbox plan test completed successfully")
+
+
+if __name__ == "__main__":
+    run_simbox_plan()

tests/integration/simbox/runners/simbox_render_runner.py

@@ -0,0 +1,29 @@
+"""
+simbox render runner for Isaac Sim subprocess execution.
+This script runs in the Isaac Sim environment and is called by the main test.
+"""
+
+import sys
+
+from tests.integration.base.test_harness import IntegrationTestHarness
+
+# Add paths to sys.path
+sys.path.append("./")
+sys.path.append("./data_engine")
+sys.path.append("./tests/integration")
+
+
+def run_simbox_render():
+    """
+    Run the simbox render test in Isaac Sim environment.
+    """
+    harness = IntegrationTestHarness(config_path="configs/simbox/de_render_template.yaml", seed=42)
+
+    # Run the pipeline and validate output
+    harness.run_test_end_to_end(min_output_files=1)
+
+    print("✓ simbox render test completed successfully")
+
+
+if __name__ == "__main__":
+    run_simbox_render()

tests/integration/simbox/test_simbox.py

@@ -0,0 +1,104 @@
+"""
+simbox tests that run in Isaac Sim environment using subprocess wrappers.
+Migrated from original test files to use the new IntegrationTestHarness framework.
+"""
+
+import sys
+
+from tests.integration.base.test_harness import IntegrationTestHarness
+
+# Add path for proper imports
+sys.path.append("./")
+sys.path.append("./data_engine")
+sys.path.append("./tests/integration")
+
+
+def test_simbox_pipeline():
+    """
+    Test simbox pipeline by running it in Isaac Sim subprocess.
+    This test uses a subprocess wrapper to handle Isaac Sim process separation.
+    """
+    harness = IntegrationTestHarness(config_path="configs/simbox/de_pipe_template.yaml", seed=42, random_num=1)
+
+    # Run in subprocess using Isaac Sim Python interpreter
+    result = harness.run_data_engine_subprocess(
+        runner_script="tests/integration/simbox/runners/simbox_pipeline_runner.py",
+        interpreter="/isaac-sim/python.sh",
+        timeout=1800,  # 30 minutes
+    )
+
+    # Verify subprocess completed successfully
+    assert result.returncode == 0, f"simbox pipeline test failed with return code: {result.returncode}"
+
+    # Validate that output was generated
+    # harness.validate_output_generated(min_files=6)
+
+    print("✓ simbox pipeline test completed successfully")
+
+
+def test_simbox_plan():
+    """
+    Test simbox plan generation by running it in Isaac Sim subprocess.
+    """
+    harness = IntegrationTestHarness(
+        config_path="configs/simbox/de_plan_template.yaml",
+        seed=42,
+        random_num=1,
+    )
+
+    # Run in subprocess using Isaac Sim Python interpreter
+    result = harness.run_data_engine_subprocess(
+        runner_script="tests/integration/simbox/runners/simbox_plan_runner.py",
+        interpreter="/isaac-sim/python.sh",
+        timeout=1800,  # 30 minutes
+    )
+
+    # Verify subprocess completed successfully
+    assert result.returncode == 0, f"simbox plan test failed with return code: {result.returncode}"
+
+    # Validate that output was generated
+    # harness.validate_output_generated(min_files=1)
+
+    print("✓ simbox plan test completed successfully")
+
+
+def test_simbox_render():
+    """
+    Test simbox render by running it in Isaac Sim subprocess.
+    """
+    harness = IntegrationTestHarness(
+        config_path="configs/simbox/de_render_template.yaml",
+        # config_path="tests/integration/simbox/configs/simbox_test_render_configs.yaml",
+        seed=42,
+    )
+
+    reference_dir = "/shared/smartbot_new/zhangyuchang/CI/manip/simbox/simbox_render_ci"
+    # Run in subprocess using Isaac Sim Python interpreter
+    result = harness.run_data_engine_subprocess(
+        runner_script="tests/integration/simbox/runners/simbox_render_runner.py",
+        interpreter="/isaac-sim/python.sh",
+        timeout=1800,  # 30 minutes
+        compare_output=True,
+        reference_dir=reference_dir,
+        comparator="simbox",
+        comparator_args={
+            # Use defaults; override here if needed
+            # "tolerance": 1e-6,
+            # "image_psnr": 37.0,
+        },
+    )
+
+    # Verify subprocess completed successfully
+    assert result.returncode == 0, f"simbox render test failed with return code: {result.returncode}"
+
+    # Validate that output was generated
+    # harness.validate_output_generated(min_files=1)
+
+    print("✓ simbox render test completed successfully")
+
+
+if __name__ == "__main__":
+    # Run all tests when script is executed directly
+    test_simbox_plan()
+    test_simbox_render()
+    test_simbox_pipeline()

tests/run_tests.sh

@@ -0,0 +1,127 @@
+#!/bin/bash
+# Integration test runner for SimBox DataEngine
+
+RED='\033[0;31m'
+GREEN='\033[0;32m'
+YELLOW='\033[1;33m'
+BLUE='\033[0;34m'
+CYAN='\033[0;36m'
+BOLD='\033[1m'
+NC='\033[0m'
+
+TIMESTAMP=$(date +%Y%m%d_%H%M%S)
+LOG_FILE="tests/test_results_${TIMESTAMP}.log"
+SUMMARY_FILE="tests/test_summary_${TIMESTAMP}.txt"
+TEMP_LOG="tests/temp_test_output.log"
+
+declare -a TEST_SUITES=()
+TEST_SUITES+=("SimBox Tests (Isaac Sim):3:/isaac-sim/python.sh tests/integration/simbox/test_simbox.py")
+
+TOTAL_SUITES=${#TEST_SUITES[@]}
+
+echo "Starting SimBox DataEngine Integration Tests..."
+echo "=============================================="
+echo -e "${BOLD}${CYAN}TEST EXECUTION PLAN:${NC}"
+echo -e "  ${BOLD}Total Test Suites: ${TOTAL_SUITES}${NC}"
+echo -e "  ${BOLD}SimBox Scenarios Covered:${NC}"
+echo -e "    - Pipeline: Full end-to-end workflow"
+echo -e "    - Plan: Trajectory planning generation"
+echo -e "    - Render: Scene rendering with validation"
+echo ""
+
+echo -e "${BLUE}Detailed logs: $LOG_FILE${NC}"
+echo "=============================================="
+
+TOTAL_TEST_SUITES=0
+PASSED_TEST_SUITES=0
+FAILED_TEST_SUITES=0
+
+run_test_suite() {
+    local suite_name="$1"
+    local expected_sessions="$2"
+    local test_command="$3"
+    local current_suite=$4
+
+    TOTAL_TEST_SUITES=$((TOTAL_TEST_SUITES + 1))
+
+    echo -e "${BOLD}${BLUE}[$current_suite/$TOTAL_SUITES] Starting: $suite_name${NC}"
+    echo -e "  ${CYAN}-> Running ${expected_sessions} SimBox tests with Isaac Sim${NC}"
+
+    echo "Test Suite: $suite_name" >> "$LOG_FILE"
+    echo "Expected Test Functions: $expected_sessions" >> "$LOG_FILE"
+    echo "Command: $test_command" >> "$LOG_FILE"
+    echo "Started at: $(date)" >> "$LOG_FILE"
+    echo "----------------------------------------" >> "$LOG_FILE"
+
+    echo -e "${CYAN}Executing: $test_command${NC}"
+    eval "$test_command" > "$TEMP_LOG" 2>&1
+    local exit_status=$?
+
+    cat "$TEMP_LOG" >> "$LOG_FILE"
+    echo "" >> "$LOG_FILE"
+
+    if [ $exit_status -eq 0 ]; then
+        echo -e "${GREEN}✓ [$current_suite/$TOTAL_SUITES] PASSED: $suite_name${NC}"
+        PASSED_TEST_SUITES=$((PASSED_TEST_SUITES + 1))
+    else
+        echo -e "${RED}✗ [$current_suite/$TOTAL_SUITES] FAILED: $suite_name${NC}"
+        FAILED_TEST_SUITES=$((FAILED_TEST_SUITES + 1))
+        echo -e "${YELLOW}  Check $LOG_FILE for error details${NC}"
+        echo -e "${YELLOW}Recent output:${NC}"
+        tail -20 "$TEMP_LOG"
+    fi
+
+    echo ""
+    rm -f "$TEMP_LOG"
+}
+
+> "$SUMMARY_FILE"
+> "$LOG_FILE"
+
+echo "Test execution started at: $(date)" > "$SUMMARY_FILE"
+echo "Planned suites: $TOTAL_SUITES" >> "$SUMMARY_FILE"
+echo "=======================================" >> "$SUMMARY_FILE"
+
+echo "Test execution started at: $(date)" >> "$LOG_FILE"
+echo "================================================" >> "$LOG_FILE"
+
+for i in "${!TEST_SUITES[@]}"; do
+    suite_info="${TEST_SUITES[$i]}"
+    suite_name=$(echo "$suite_info" | cut -d':' -f1)
+    expected_sessions=$(echo "$suite_info" | cut -d':' -f2)
+    test_command=$(echo "$suite_info" | cut -d':' -f3-)
+    current_suite=$((i + 1))
+
+    run_test_suite "$suite_name" "$expected_sessions" "$test_command" $current_suite
+done
+
+echo "=============================================="
+echo "TEST EXECUTION SUMMARY"
+echo "=============================================="
+echo -e "${CYAN}Test Suites:${NC}"
+echo -e "  Total: $TOTAL_TEST_SUITES"
+echo -e "  ${GREEN}Passed: $PASSED_TEST_SUITES${NC}"
+echo -e "  ${RED}Failed: $FAILED_TEST_SUITES${NC}"
+echo ""
+
+echo "" >> "$SUMMARY_FILE"
+echo "FINAL SUMMARY:" >> "$SUMMARY_FILE"
+echo "Suites - Total: $TOTAL_TEST_SUITES, Passed: $PASSED_TEST_SUITES, Failed: $FAILED_TEST_SUITES" >> "$SUMMARY_FILE"
+
+if [ $FAILED_TEST_SUITES -eq 0 ]; then
+    echo -e "${GREEN}ALL TESTS PASSED${NC}"
+    echo "ALL TESTS PASSED at $(date)" >> "$SUMMARY_FILE"
+else
+    echo -e "${RED}SOME TEST SUITES FAILED${NC}"
+    echo "SOME TEST SUITES FAILED at $(date)" >> "$SUMMARY_FILE"
+fi
+
+echo ""
+echo -e "${BLUE}Results: $SUMMARY_FILE${NC}"
+echo -e "${BLUE}Logs: $LOG_FILE${NC}"
+
+if [ $FAILED_TEST_SUITES -eq 0 ]; then
+    exit 0
+else
+    exit 1
+fi