base_generator.py

from abc import ABC, abstractmethod
from typing import Any, ClassVar, Dict, List, Optional, Type, TypeVar, Generic

import rospy
from rosnav_rl.states import SimulationStateContainer

from ..collectors import BaseUnit

GeneratedDataType = TypeVar("D")


class GenerationError(Exception):
    """Exception raised when there's an error generating an observation."""
    
    def __init__(self, generator_name: str, message: str, original_exception: Optional[Exception] = None) -> None:
        self.generator_name = generator_name
        self.original_exception = original_exception
        super().__init__(f"Error in generator '{generator_name}': {message}")


class ObservationGeneratorUnit(BaseUnit, Generic[GeneratedDataType], ABC):
    """
    Base class for observation generator units.

    This class serves as a foundation for units that generate processed observations
    from other observation data. Generators can have dependencies on other observation
    units, which are verified before attempting to generate data.

    Class Attributes:
        name (ClassVar[str]): The name identifier for this generator.
        requires (ClassVar[List[BaseUnit]]): A list of required base units.
        data_class (Type[GeneratedDataType]): The type of data generated by this unit.
        fallback_value (ClassVar[Optional[Any]]): Fallback value to use when generation fails.
        max_consecutive_errors (ClassVar[int]): Maximum allowed consecutive errors before failing.
    
    Instance Attributes:
        _error_count (int): Count of total errors encountered.
        _consecutive_errors (int): Count of consecutive errors without successful generation.
        
    Args:
        *args: Variable length argument list.
        **kwargs: Arbitrary keyword arguments.
    """

    name: ClassVar[str]
    requires: ClassVar[List[BaseUnit]]
    data_class: Type[GeneratedDataType]
    
    fallback_value: ClassVar[Optional[Any]] = None
    max_consecutive_errors: ClassVar[int] = 3

    def __init__(self, *args, **kwargs) -> None:
        """Initialize the observation generator unit."""
        super().__init__(*args, **kwargs)
        self._error_count = 0
        self._consecutive_errors = 0

    def validate_requirements(self, obs_dict: Dict[str, Any]) -> List[str]:
        """
        Check if all required observations are present in the observation dictionary.
        
        Args:
            obs_dict (Dict[str, Any]): Dictionary containing observations
            
        Returns:
            List[str]: List of missing requirements (empty if all are present)
        """
        return [
            req.name for req in self.requires 
            if req.name not in obs_dict or obs_dict[req.name] is None
        ]

    def validate_input(self, obs_dict: Dict[str, Any]) -> bool:
        """
        Validate input observations beyond checking for presence.
        
        This can be overridden by subclasses to perform additional validation
        of input observations.
        
        Args:
            obs_dict (Dict[str, Any]): Dictionary containing observations
            
        Returns:
            bool: True if inputs are valid, False otherwise
        """
        return True

    def safe_generate(
        self,
        obs_dict: Dict[str, Any],
        simulation_state_container: SimulationStateContainer,
        *args,
        **kwargs,
    ) -> GeneratedDataType:
        """
        Safely generate observation data with error handling.
        
        This method wraps the generate method with proper error handling,
        dependency checking, and fallback mechanisms.
        
        Args:
            obs_dict (Dict[str, Any]): Dictionary containing observations
            simulation_state_container (SimulationStateContainer): Container for simulation state
            *args: Additional positional arguments
            **kwargs: Additional keyword arguments
            
        Returns:
            GeneratedDataType: The generated observation data
            
        Raises:
            GenerationError: If generation fails and error tolerance is exceeded
        """
        # Check requirements
        missing_requirements = self.validate_requirements(obs_dict)
        if missing_requirements:
            error_msg = f"Missing required observations: {', '.join(missing_requirements)}"
            rospy.logwarn_once(f"[{self.name}] {error_msg}")
            
            self._error_count += 1
            self._consecutive_errors += 1
            
            if self._consecutive_errors > self.max_consecutive_errors:
                raise GenerationError(
                    generator_name=self.name,
                    message=error_msg
                )
            if self.fallback_value is not None:
                return self.fallback_value
            raise GenerationError(
                generator_name=self.name,
                message=error_msg
            )
        
        # Validate inputs
        if not self.validate_input(obs_dict):
            error_msg = "Input validation failed"
            rospy.logwarn_once(f"[{self.name}] {error_msg}")
            
            self._error_count += 1
            self._consecutive_errors += 1
            
            if self._consecutive_errors > self.max_consecutive_errors:
                raise GenerationError(
                    generator_name=self.name,
                    message=error_msg
                )
            if self.fallback_value is not None:
                return self.fallback_value
            raise GenerationError(
                generator_name=self.name,
                message=error_msg
            )
        
        # Generate observation
        try:
            result = self.generate(
                obs_dict=obs_dict,
                simulation_state_container=simulation_state_container,
                *args, 
                **kwargs
            )
            # Reset consecutive errors on success
            self._consecutive_errors = 0
            return result
        except Exception as e:
            error_msg = f"Generation failed: {str(e)}"
            rospy.logwarn_throttle(5.0, f"[{self.name}] {error_msg}")
            
            self._error_count += 1
            self._consecutive_errors += 1
            
            if self._consecutive_errors > self.max_consecutive_errors:
                raise GenerationError(
                    generator_name=self.name,
                    message=error_msg,
                    original_exception=e
                )
            if self.fallback_value is not None:
                return self.fallback_value
            raise GenerationError(
                generator_name=self.name,
                message=error_msg,
                original_exception=e
            )

    @abstractmethod
    def generate(
        self,
        obs_dict: Dict[str, Any],
        simulation_state_container: SimulationStateContainer,
        *args,
        **kwargs,
    ) -> GeneratedDataType:
        """
        Generates the observation data based on the given observation dictionary.
        
        Implementations should focus on the core generation logic without worrying
        about error handling or dependency checking, which is handled by safe_generate.

        Args:
            obs_dict (Dict[str, Any]): The observation dictionary.
            simulation_state_container (SimulationStateContainer): The simulation state container.
            *args: Additional positional arguments.
            **kwargs: Additional keyword arguments.

        Returns:
            GeneratedDataType: The generated observation data.
            
        Raises:
            NotImplementedError: If the method is not implemented by a subclass.
        """
        raise NotImplementedError