ding.framework.middleware.functional.priority

ding.framework.middleware.functional.priority

priority_calculator(priority_calculation_fn)

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.

Full Source Code

../ding/framework/middleware/functional/priority.py

1from typing import TYPE_CHECKING, Callable 2from ding.framework import task 3if TYPE_CHECKING: 4 from ding.framework import OnlineRLContext 5 6 7def priority_calculator(priority_calculation_fn: Callable) -> Callable: 8 """ 9 Overview: 10 The middleware that calculates the priority of the collected data. 11 Arguments: 12 - priority_calculation_fn (:obj:`Callable`): The function that calculates the priority of the collected data. 13 """ 14 15 if task.router.is_active and not task.has_role(task.role.COLLECTOR): 16 return task.void() 17 18 def _priority_calculator(ctx: "OnlineRLContext") -> None: 19 20 priority = priority_calculation_fn(ctx.trajectories) 21 for i in range(len(priority)): 22 ctx.trajectories[i]['priority'] = priority[i] 23 24 return _priority_calculator