| from typing import TYPE_CHECKING, Callable | |
| from ding.framework import task | |
| if TYPE_CHECKING: | |
| from ding.framework import OnlineRLContext | |
| def priority_calculator(priority_calculation_fn: Callable) -> Callable: | |
| """ | |
| Overview: | |
| The middleware that calculates the priority of the collected data. | |
| Arguments: | |
| - priority_calculation_fn (:obj:`Callable`): The function that calculates the priority of the collected data. | |
| """ | |
| if task.router.is_active and not task.has_role(task.role.COLLECTOR): | |
| return task.void() | |
| def _priority_calculator(ctx: "OnlineRLContext") -> None: | |
| priority = priority_calculation_fn(ctx.trajectories) | |
| for i in range(len(priority)): | |
| ctx.trajectories[i]['priority'] = priority[i] | |
| return _priority_calculator | |