DI-engine/ding/data/buffer/middleware/use_time_check.py

from typing import Callable, Any, List, Optional, Union, TYPE_CHECKING
from collections import defaultdict
from ding.data.buffer import BufferedData
if TYPE_CHECKING:
    from ding.data.buffer.buffer import Buffer


def use_time_check(buffer_: 'Buffer', max_use: int = float("inf")) -> Callable:
    """
    Overview:
        This middleware aims to check the usage times of data in buffer. If the usage times of a data is
        greater than or equal to max_use, this data will be removed from buffer as soon as possible.
    Arguments:
        - max_use (:obj:`int`): The max reused (resampled) count for any individual object.
    """

    use_count = defaultdict(int)

    def _need_delete(item: BufferedData) -> bool:
        nonlocal use_count
        idx = item.index
        use_count[idx] += 1
        item.meta['use_count'] = use_count[idx]
        if use_count[idx] >= max_use:
            return True
        else:
            return False

    def _check_use_count(sampled_data: List[BufferedData]):
        delete_indices = [item.index for item in filter(_need_delete, sampled_data)]
        buffer_.delete(delete_indices)
        for index in delete_indices:
            del use_count[index]

    def sample(chain: Callable, *args, **kwargs) -> Union[List[BufferedData], List[List[BufferedData]]]:
        sampled_data = chain(*args, **kwargs)
        if len(sampled_data) == 0:
            return sampled_data

        if isinstance(sampled_data[0], BufferedData):
            _check_use_count(sampled_data)
        else:
            for grouped_data in sampled_data:
                _check_use_count(grouped_data)
        return sampled_data

    def _use_time_check(action: str, chain: Callable, *args, **kwargs) -> Any:
        if action == "sample":
            return sample(chain, *args, **kwargs)
        return chain(*args, **kwargs)

    return _use_time_check